Update BidirectionalTranslation/models/networks.py

BidirectionalTranslation/models/networks.py

@@ -57,8 +57,7 @@ def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[], init=True):
         gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
     Return an initialized network.
     """
-    if len(gpu_ids) > 0:
-        assert(torch.cuda.is_available())
+    if len(gpu_ids) > 0 and torch.cuda.is_available():
         net.to(gpu_ids[0])
     if init:
         init_weights(net, init_type, init_gain=init_gain)
@@ -97,7 +96,7 @@ class LayerNormWarpper(nn.Module):
         self.num_features = int(num_features)
 
     def forward(self, x):
-        x = nn.LayerNorm([self.num_features, x.size()[2], x.size()[3]], elementwise_affine=False).cuda()(x)
+        x = nn.LayerNorm([self.num_features, x.size()[2], x.size()[3]], elementwise_affine=False).to(x.device)(x)
         return x
 
 def get_norm_layer(norm_type='instance'):
@@ -904,7 +903,7 @@ class G_Unet_add_all(nn.Module):
         for layer_idx in range(num_layers):
             res = layer_idx // 2 + 2
             shape = [1, 1, 2 ** res, 2 ** res]
-            self.noise_inputs.append(torch.randn(*shape).to("cuda"))
+            self.noise_inputs.append(torch.randn(*shape).to("cuda" if torch.cuda.is_available() else "cpu"))
 
         # construct unet structure
         unet_block = UnetBlock_with_z(ngf * 8, ngf * 8, ngf * 8, nz, submodule=None, innermost=True, 
@@ -1338,7 +1337,7 @@ class ScreenVAE(nn.Module):
                     net = net.module
                 print('loading the model from %s' % load_path)
                 state_dict = torch.load(
-                    load_path, map_location=lambda storage, loc: storage.cuda())
+                    load_path, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
                 if hasattr(state_dict, '_metadata'):
                     del state_dict._metadata
 
@@ -1372,4 +1371,4 @@ class ScreenVAE(nn.Module):
             x = self.npad(x)
             recons = self.dec(x)[:,:,:h,:w]
             recons = (recons+1)*(line+1)/2-1
-            return torch.clamp(recons,-1,1)
+            return torch.clamp(recons,-1,1)