lychiyu/freqtrade_origin
1
0
Fork 0
mirror of https://github.com/freqtrade/freqtrade.git synced 2024-09-20 09:31:12 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

pytorch - trainer - add device arg to load method

Browse Source
This commit is contained in:
yinon 2023-07-13 15:39:47 +00:00
parent 0c9aa86885
commit 49a7de4ebd
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
freqtrade/freqai/torch/PyTorchModelTrainer.py
View File

@ -182,8 +182,8 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
"pytrainer": self "pytrainer": self
}, path) }, path)
def load(self, path: Path): def load(self, path: Path, device: str = None):
checkpoint = torch.load(path) checkpoint = torch.load(path, map_location=device)
return self.load_from_checkpoint(checkpoint) return self.load_from_checkpoint(checkpoint)
def load_from_checkpoint(self, checkpoint: Dict): def load_from_checkpoint(self, checkpoint: Dict):