liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-17 17:34:35 +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

2
qlib/contrib/model/pytorch_gats_ts.py
View File

@@ -276,7 +276,7 @@ class GATs(Model):
if self.model_path is not None:
self.logger.info("Loading pretrained model...")
pretrained_model.load_state_dict(torch.load(self.model_path))
pretrained_model.load_state_dict(torch.load(self.model_path, map_location=self.device))
model_dict = self.GAT_model.state_dict()
pretrained_dict = {k: v for k, v in pretrained_model.state_dict().items() if k in model_dict}