Support for MPS device for mixed-precision

laclouis5 · laclouis5 · commit 92b13a8aeccb · 2025-01-06T19:57:40.000+01:00
diff --git a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py
@@ -520,7 +520,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 = "cpu" if self._accelerator_flag == "cpu" else "mps" if self._accelerator_flag == "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
`@@ -520,7 +520,7 @@ def _check_and_init_precision(self) -> Precision:`
`520`	`520`	`rank_zero_info(`
`521`	`521`	`f"Using {'16bit' if self._precision_flag == '16-mixed' else 'bfloat16'} Automatic Mixed Precision (AMP)"`
`522`	`522`	`)`
`523`		`- device = "cpu" if self._accelerator_flag == "cpu" else "cuda"`
	`523`	`+ device = "cpu" if self._accelerator_flag == "cpu" else "mps" if self._accelerator_flag == "mps" else "cuda"`
`524`	`524`	`return MixedPrecision(self._precision_flag, device) # type: ignore[arg-type]`
`525`	`525`
`526`	`526`	`raise RuntimeError("No precision set")`