lychiyu/freqtrade_origin
1
0
Fork 0
mirror of https://github.com/freqtrade/freqtrade.git synced 2024-11-10 10:21:59 +00:00
Code Issues Actions 2 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
}, path)
def load(self, path: Path):
checkpoint = torch.load(path)
def load(self, path: Path, device: str = None):
checkpoint = torch.load(path, map_location=device)
return self.load_from_checkpoint(checkpoint)
def load_from_checkpoint(self, checkpoint: Dict):