diff options
| author | LambdaWill <574819595@qq.com> | 2017-10-24 12:21:00 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2017-10-24 12:21:00 +0800 |
| commit | 343dda259b4b6a64a338fa61a9f1c70893b8fc7e (patch) | |
| tree | a7c98bca3224dd2457234ba7701b5b32064127ff /models/base_model.py | |
| parent | 9d0295bbb5d01ce40bd2792f72882128282f0338 (diff) | |
Update base_model.py
Diffstat (limited to 'models/base_model.py')
| -rw-r--r-- | models/base_model.py | 2 |
1 files changed, 1 insertions, 1 deletions
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): |