From 6e4fcb238cedf039951c1f5db6a3c83f36ce1d94 Mon Sep 17 00:00:00 2001 From: aceniccola Date: Sat, 23 Aug 2025 12:09:25 -0400 Subject: [PATCH] mps support --- src/state/_cli/_tx/_train.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/state/_cli/_tx/_train.py b/src/state/_cli/_tx/_train.py index 7be3afff..da3abdb6 100644 --- a/src/state/_cli/_tx/_train.py +++ b/src/state/_cli/_tx/_train.py @@ -213,6 +213,8 @@ def run_tx_train(cfg: DictConfig): if torch.cuda.is_available(): accelerator = "gpu" + elif torch.backends.mps.is_available(): + accelerator = "mps" else: accelerator = "cpu" @@ -259,7 +261,7 @@ def run_tx_train(cfg: DictConfig): if checkpoint_path is None and manual_init is not None: print(f"Loading manual checkpoint from {manual_init}") checkpoint_path = manual_init - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu") checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) model_state = model.state_dict() checkpoint_state = checkpoint["state_dict"]