Update BidirectionalTranslation/models/cycle_ganstft_model.py

BidirectionalTranslation/models/cycle_ganstft_model.py

@@ -16,8 +16,8 @@ class CycleGANSTFTModel(BaseModel):
         self.interchnnls = 4
         use_noise = False
         self.half_size = opt.batch_size //2
-        self.device=opt.local_rank
-        self.gpu_ids=[self.device]
+        self.device = torch.device('cuda:{}'.format(opt.local_rank)) if torch.cuda.is_available() else torch.device('cpu')
+        self.gpu_ids = [opt.local_rank] if torch.cuda.is_available() else []
         self.local_rank = opt.local_rank
         self.cropsize = opt.crop_size
 
@@ -100,4 +100,4 @@ class CycleGANSTFTModel(BaseModel):
                     k += 1
                 z = torch.clamp(z, -tvalue, tvalue)
 
-        return z.detach().to(self.device)
+        return z.detach().to(self.device)