[Reland] Add warning about removed sm50 and sm60 arches

atalman · atalman · commit 5c616e2c33c9 · 2025-07-21T08:12:30.000-07:00
diff --git a/torch/cuda/__init__.py b/torch/cuda/__init__.py
@@ -262,7 +262,9 @@ def _check_capability():
         "12.9": {"min": 70, "max": 120},
     }
 
-    if torch.version.cuda is not None:  # on ROCm we don't want this check
+    if (
+        torch.version.cuda is not None and torch.cuda.get_arch_list()
+    ):  # on ROCm we don't want this check
         for d in range(device_count()):
             capability = get_device_capability(d)
             major = capability[0]
