From 455e803e6b41580297e5c4a9f192a1ba891477c3 Mon Sep 17 00:00:00 2001 From: junyanz Date: Sat, 10 Feb 2018 19:57:10 -0500 Subject: fix cuda issue --- models/cycle_gan_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'models/cycle_gan_model.py') diff --git a/models/cycle_gan_model.py b/models/cycle_gan_model.py index bd21d28..730d077 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.cuda(self.gpu_ids[0], async=True) - input_B.cuda(self.gpu_ids[0], async=True) + input_A = input_A.cuda(self.gpu_ids[0], async=True) + input_B = 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'] -- cgit v1.2.3-70-g09d2