Merge pull request #7 from abhilash1910/patch-5

Fix precision device
Lightning-AI · Oct 31, 2023 · 5bfb98a · 5bfb98a
2 parents 1b90607 + 893b654
commit 5bfb98a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/lightning/pytorch/trainer/connectors/accelerator_connector.py b/src/lightning/pytorch/trainer/connectors/accelerator_connector.py
@@ -589,7 +589,7 @@ def _check_and_init_precision(self) -> PrecisionPlugin:
             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 "xpu" if _lightning_xpu_available() else "cuda"
             return MixedPrecisionPlugin(self._precision_flag, device)  # type: ignore[arg-type]
 
         raise RuntimeError("No precision set")