gpu tensor to cpu

mrnet/train.py

@@ -53,8 +53,7 @@ def train_net(net, device, epochs = 5, batch_size = 1, lr = 0.1):
         jac = eval_jac(net,val_loader,device,n_val)
         # overall_acc, avg_per_class_acc, avg_jacc, avg_dice = eval_multi(net, val_loader, device, n_val)
         scheduler.step()
-        logging.info(f'Avg Dice:{dice}\n'
-                     f'Jaccard:{jac}\n'
+        logging.info(f'Avg Dice:{dice}   Jaccard:{jac}\n'
                      f'Learning Rate:{scheduler.get_lr()[0]}')
         if epoch % 5 == 0:
             try:

utils/eval.py

@@ -64,8 +64,8 @@ def eval_jac(net, loader, device, n_val):
             true_masks = true_masks.to(device = device, dtype = mask_type)
             pred_masks = net(imgs)
 
-            pred_masks = torch.round(pred_masks).detach().numpy()
-            true_masks = torch.round(true_masks).numpy()
+            pred_masks = torch.round(pred_masks).cpu().detach().numpy()
+            true_masks = torch.round(true_masks).cpu().numpy()
             jac += jaccard_score(true_masks.flatten(), pred_masks.flatten())
 
             pbar.update(imgs.shape[0])