问题遇到的现象和发生背景
服务器跑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()
