update amp

zhiyuan1i · zhiyuan1i · commit 06a3303742ec · 2024-12-20T03:20:18.000Z
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":

Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,7 @@ def __init__(`
`56`	`56`	`if _TORCH_GREATER_EQUAL_2_4`
`57`	`57`	`else getattr(`
`58`	`58`	`torch,`
`59`		`- "cuda" if not isinstance(device, str) or device.split(":")[0] == "cpu" else device.split(":")[0],`
	`59`	`+ "cuda" if device.split(":")[0] == "cpu" else device.split(":")[0],`
`60`	`60`	`).amp.GradScaler()`
`61`	`61`	`)`
`62`	`62`	`if scaler is not None and self.precision == "bf16-mixed":`