From 49a7de4ebdf7ea3621f072957ca936b26f82fc03 Mon Sep 17 00:00:00 2001 From: yinon Date: Thu, 13 Jul 2023 15:39:47 +0000 Subject: [PATCH] pytorch - trainer - add device arg to load method --- freqtrade/freqai/torch/PyTorchModelTrainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/freqtrade/freqai/torch/PyTorchModelTrainer.py b/freqtrade/freqai/torch/PyTorchModelTrainer.py index e74b572fd..b49e16196 100644 --- a/freqtrade/freqai/torch/PyTorchModelTrainer.py +++ b/freqtrade/freqai/torch/PyTorchModelTrainer.py @@ -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):