update amp

Lightning-AI · Dec 20, 2024 · 06a3303 · 06a3303
1 parent 5cdd9e7
commit 06a3303
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/lightning/fabric/plugins/precision/amp.py b/src/lightning/fabric/plugins/precision/amp.py
@@ -56,7 +56,7 @@ def __init__(
                 if _TORCH_GREATER_EQUAL_2_4
                 else getattr(
                     torch,
-                    "cuda" if not isinstance(device, str) or device.split(":")[0] == "cpu" else device.split(":")[0],
+                    "cuda" if device.split(":")[0] == "cpu" else device.split(":")[0],
                 ).amp.GradScaler()
             )
         if scaler is not None and self.precision == "bf16-mixed":

diff --git a/src/lightning/pytorch/plugins/precision/amp.py b/src/lightning/pytorch/plugins/precision/amp.py
@@ -56,7 +56,7 @@ def __init__(
                 if _TORCH_GREATER_EQUAL_2_4
                 else getattr(
                     torch,
-                    "cuda" if not isinstance(device, str) or device.split(":")[0] == "cpu" else device.split(":")[0],
+                    "cuda" if device.split(":")[0] == "cpu" else device.split(":")[0],
                 ).amp.GradScaler()
             )
         if scaler is not None and self.precision == "bf16-mixed":