Merge pull request vqdang#109 from alfrei/fix-batch-size-1-error

torch.squeeze causes an error for valid batch size 1
samjmolyneux · Apr 6, 2021 · abf6d89 · abf6d89
2 parents be8ae2d + 0306601
commit abf6d89
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/models/hovernet/run_desc.py b/models/hovernet/run_desc.py
@@ -35,8 +35,8 @@ def train_step(batch_data, run_info):
     imgs = imgs.permute(0, 3, 1, 2).contiguous()
 
     # HWC
-    true_np = torch.squeeze(true_np).to("cuda").type(torch.int64)
-    true_hv = torch.squeeze(true_hv).to("cuda").type(torch.float32)
+    true_np = true_np.to("cuda").type(torch.int64)
+    true_hv = true_hv.to("cuda").type(torch.float32)
 
     true_np_onehot = (F.one_hot(true_np, num_classes=2)).type(torch.float32)
     true_dict = {