From 343dda259b4b6a64a338fa61a9f1c70893b8fc7e Mon Sep 17 00:00:00 2001
From: LambdaWill <574819595@qq.com>
Date: Tue, 24 Oct 2017 12:21:00 +0800
Subject: Update base_model.py

---
 models/base_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'models/base_model.py')

diff --git a/models/base_model.py b/models/base_model.py
index 446a903..d62d189 100644
--- a/models/base_model.py
+++ b/models/base_model.py
@@ -44,7 +44,7 @@ class BaseModel():
         save_path = os.path.join(self.save_dir, save_filename)
         torch.save(network.cpu().state_dict(), save_path)
         if len(gpu_ids) and torch.cuda.is_available():
-            network.cuda(device_id=gpu_ids[0])
+            network.cuda(device_id=gpu_ids[0]) # network.cuda(device=gpu_ids[0]) for the latest version.
 
     # helper loading function that can be used by subclasses
     def load_network(self, network, network_label, epoch_label):
-- 
cgit v1.2.3-70-g09d2