summaryrefslogtreecommitdiff
path: root/models/test_model.py
diff options
context:
space:
mode:
Diffstat (limited to 'models/test_model.py')
-rw-r--r--models/test_model.py2
1 files changed, 1 insertions, 1 deletions
diff --git a/models/test_model.py b/models/test_model.py
index 35b7402..f593c46 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.cuda(self.gpu_ids[0], async=True)
+ input_A = input_A.cuda(self.gpu_ids[0], async=True)
self.input_A = input_A
self.image_paths = input['A_paths']