Update device assignment logic to support 'mps' accelerator

bhimrazy · bhimrazy · commit 92f01038c4bf · 2025-09-14T23:39:48.000+05:45
diff --git a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py
@@ -515,7 +515,7 @@ def _check_and_init_precision(self) -> Precision:
             rank_zero_info(
                 f"Using {'16bit' if self._precision_flag == '16-mixed' else 'bfloat16'} Automatic Mixed Precision (AMP)"
             )
-            device = "cpu" if self._accelerator_flag == "cpu" else "cuda"
+            device = self._accelerator_flag if self._accelerator_flag in ("cpu", "mps") else "cuda"
             return MixedPrecision(self._precision_flag, device)  # type: ignore[arg-type]
 
         raise RuntimeError("No precision set")

Original file line number	Diff line number	Diff line change
`@@ -515,7 +515,7 @@ def _check_and_init_precision(self) -> Precision:`
`515`	`515`	`rank_zero_info(`
`516`	`516`	`f"Using {'16bit' if self._precision_flag == '16-mixed' else 'bfloat16'} Automatic Mixed Precision (AMP)"`
`517`	`517`	`)`
`518`		`- device = "cpu" if self._accelerator_flag == "cpu" else "cuda"`
	`518`	`+ device = self._accelerator_flag if self._accelerator_flag in ("cpu", "mps") else "cuda"`
`519`	`519`	`return MixedPrecision(self._precision_flag, device) # type: ignore[arg-type]`
`520`	`520`
`521`	`521`	`raise RuntimeError("No precision set")`