diff --git a/freqtrade/freqai/torch/PyTorchModelTrainer.py b/freqtrade/freqai/torch/PyTorchModelTrainer.py
index e6691f3db..e6638d4fd 100644
--- a/freqtrade/freqai/torch/PyTorchModelTrainer.py
+++ b/freqtrade/freqai/torch/PyTorchModelTrainer.py
@@ -176,7 +176,7 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
             "pytrainer": self
         }, path)
 
-    def load(self, path: Path, device: str = None):
+    def load(self, path: Path, device: Optional[str] = None):
         checkpoint = torch.load(path, map_location=device)
         return self.load_from_checkpoint(checkpoint)