liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-16 01:06:56 +08:00
Code Activity

add map_location to torch.load to make it work when cuda is unavailable (#782)

Browse Source
This commit is contained in:
Chao Ning
2021-12-28 16:02:04 +00:00
committed by GitHub
parent 6bafd0a09b
commit 622303b83a
5 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
qlib/contrib/model/pytorch_nn.py
View File

@@ -257,7 +257,7 @@ class DNNModelPytorch(Model):
self.scheduler.step(cur_loss_val)
# restore the optimal parameters after training
self.dnn_model.load_state_dict(torch.load(save_path))
self.dnn_model.load_state_dict(torch.load(save_path, map_location=self.device))
if self.use_gpu:
torch.cuda.empty_cache()
@@ -296,7 +296,7 @@ class DNNModelPytorch(Model):
]
_model_path = os.path.join(model_dir, _model_name)
# Load model
self.dnn_model.load_state_dict(torch.load(_model_path))
self.dnn_model.load_state_dict(torch.load(_model_path, map_location=self.device))
self.fitted = True