Merge pull request #10773 from freqtrade/fix/freqai_m3

fix: Support mps device where available
This commit is contained in:
Matthias 2024-10-18 07:07:23 +02:00 committed by GitHub
commit e47042d1cb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -20,7 +20,11 @@ class BasePyTorchModel(IFreqaiModel, ABC):
def __init__(self, **kwargs):
super().__init__(config=kwargs["config"])
self.dd.model_type = "pytorch"
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.device = (
"mps"
if torch.backends.mps.is_available() and torch.backends.mps.is_built()
else ("cuda" if torch.cuda.is_available() else "cpu")
)
test_size = self.freqai_info.get("data_split_parameters", {}).get("test_size")
self.splits = ["train", "test"] if test_size != 0 else ["train"]
self.window_size = self.freqai_info.get("conv_width", 1)