TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu()

Question

问题遇到的现象和发生背景

服务器跑train.py

问题相关代码

        for val_img, val_mask in tqdm(val_data_loader_num, ncols=20, total=len(val_data_loader_num)):
            val_img, val_mask = val_img.to(device), val_mask.cuda()
            val_mask[np.where(val_mask > 0)] = 1
            val_mask = val_mask.squeeze(0)
            predict = solver.test_one_img(val_img)
            predict_temp = torch.from_numpy(predict).unsqueeze(0)
            predict_use = V(predict_temp.type(torch.FloatTensor), volatile=True)
            val_use = V(val_mask.type(torch.FloatTensor), volatile=True)
            test_epoch_loss += criteon.forward(predict_use, val_use)
            predict_use = predict_use.squeeze(0)
            predict_use = predict_use.unsqueeze(1)
            predict_use[predict_use >= 0.5] = 1
            predict_use[predict_use < 0.5] = 0
            predict_use = predict_use.type(torch.LongTensor)
            val_use = val_use.squeeze(1).type(torch.LongTensor)
            test_mean_iou += iou_pytorch(predict_use, val_use)

运行结果及报错内容

---------- Epoch:1 ----------
Train:
100%|?| 463/463 [10:
Validation:
  0%| | 0/66 [00:02<
Traceback (most recent call last):
  File "train.py", line 239, in <module>
    val_mask[np.where(val_mask > 0)] = 1
  File "/public/home/G19860004/anaconda3/lib/python3.6/site-packages/torch/tensor.py", line 458, in __array__
    return self.numpy()
TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

我的解答思路和尝试过的方法

尝试改为val_mask[np.where(val_mask > 0).cpu()] = 1，依旧报错，不知道该从哪里加.cpu()

Answer 1

在val_mask.cuda()后面加了.cpu() 就可以啦