diff options
| -rw-r--r-- | models/cycle_gan_model.py | 4 | ||||
| -rw-r--r-- | models/pix2pix_model.py | 4 | ||||
| -rw-r--r-- | models/test_model.py | 2 |
3 files changed, 5 insertions, 5 deletions
diff --git a/models/cycle_gan_model.py b/models/cycle_gan_model.py index 730d077..bd21d28 100644 --- a/models/cycle_gan_model.py +++ b/models/cycle_gan_model.py @@ -72,8 +72,8 @@ class CycleGANModel(BaseModel): input_A = input['A' if AtoB else 'B'] input_B = input['B' if AtoB else 'A'] if len(self.gpu_ids) > 0: - input_A = input_A.cuda(self.gpu_ids[0], async=True) - input_B = input_B.cuda(self.gpu_ids[0], async=True) + input_A.cuda(self.gpu_ids[0], async=True) + input_B.cuda(self.gpu_ids[0], async=True) self.input_A = input_A self.input_B = input_B self.image_paths = input['A_paths' if AtoB else 'B_paths'] diff --git a/models/pix2pix_model.py b/models/pix2pix_model.py index 78f8d69..79c3fd7 100644 --- a/models/pix2pix_model.py +++ b/models/pix2pix_model.py @@ -57,8 +57,8 @@ class Pix2PixModel(BaseModel): input_A = input['A' if AtoB else 'B'] input_B = input['B' if AtoB else 'A'] if len(self.gpu_ids) > 0: - input_A = input_A.cuda(self.gpu_ids[0], async=True) - input_B = input_B.cuda(self.gpu_ids[0], async=True) + input_A.cuda(self.gpu_ids[0], async=True) + input_B.cuda(self.gpu_ids[0], async=True) self.input_A = input_A self.input_B = input_B self.image_paths = input['A_paths' if AtoB else 'B_paths'] diff --git a/models/test_model.py b/models/test_model.py index f593c46..35b7402 100644 --- a/models/test_model.py +++ b/models/test_model.py @@ -28,7 +28,7 @@ class TestModel(BaseModel): # we need to use single_dataset mode input_A = input['A'] if len(self.gpu_ids) > 0: - input_A = input_A.cuda(self.gpu_ids[0], async=True) + input_A.cuda(self.gpu_ids[0], async=True) self.input_A = input_A self.image_paths = input['A_paths'] |