Update BidirectionalTranslation/models/base_model.py

BidirectionalTranslation/models/base_model.py

@@ -36,7 +36,7 @@ class BaseModel(ABC):
         self.iter = 0
         self.last_iter = 0
         self.device = torch.device('cuda:{}'.format(
-            self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')  # get device name: CPU or GPU
+            self.gpu_ids[0])) if self.gpu_ids and torch.cuda.is_available() else torch.device('cpu')  # get device name: CPU or GPU
         # save all the checkpoints to save_dir
         self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
         try:
@@ -231,7 +231,7 @@ class BaseModel(ABC):
                 # if you are using PyTorch newer than 0.4 (e.g., built from
                 # GitHub source), you can remove str() on self.device
                 state_dict = torch.load(
-                    load_path, map_location=lambda storage, loc: storage.cuda())
+                    load_path, map_location=self.device)
                 if hasattr(state_dict, '_metadata'):
                     del state_dict._metadata
 
@@ -274,4 +274,4 @@ class BaseModel(ABC):
         for net in nets:
             if net is not None:
                 for param in net.parameters():
-                    param.requires_grad = requires_grad
+                    param.requires_grad = requires_grad