Upload net34.py

nets/net34.py

@@ -194,6 +194,7 @@ class Net(nn.Module):
 
     def get_fmaps(self, images_, B, T, sw, is_training, nograd_backbone):
         _, _, H_pad, W_pad = images_.shape # revised HW
+        device = images.device
 
         C, H8, W8 = self.dim*2, H_pad//8, W_pad//8
         if self.no_split:
@@ -205,7 +206,7 @@ class Net(nn.Module):
             fmaps = []
             for t in range(0, T, fmaps_chunk_size):
                 images_chunk = images[:, t : t + fmaps_chunk_size]
-                images_chunk = images_chunk.cuda()
+                images_chunk = images_chunk.to(device)
                 if self.use_basicencoder:
                     if self.full_split:
                         fmaps_chunk1 = self.fnet(images_chunk.reshape(-1, 3, H_pad, W_pad))
@@ -224,7 +225,7 @@ class Net(nn.Module):
         else:
             if not is_training:
                 # sometimes we need to move things to cuda here
-                images_ = images_.cuda()
+                images_ = images_.to(device)
             if self.use_basicencoder:
                 if self.full_split:
                     # if self.half_corr: