ROCm: don't warn about CUDA compute capabilities (pytorch#35949)

Summary:
The warning doesn't apply for ROCm.
Pull Request resolved: pytorch#35949

Differential Revision: D21050540

Pulled By: ezyang

fbshipit-source-id: 23b13bddd13f132c2017ddea12b2dc54f611fba6

Author: t-vi

torch/cuda/__init__.py

@@ -77,16 +77,17 @@ def _check_capability():
     The minimum cuda capability that we support is 3.5.
     """
 
-    CUDA_VERSION = torch._C._cuda_getCompiledVersion()
-    for d in range(device_count()):
-        capability = get_device_capability(d)
-        major = capability[0]
-        minor = capability[1]
-        name = get_device_name(d)
-        if capability == (3, 0) or major < 3:
-            warnings.warn(old_gpu_warn % (d, name, major, capability[1]))
-        elif CUDA_VERSION <= 9000 and major >= 7 and minor >= 5:
-            warnings.warn(incorrect_binary_warn % (d, name, 10000, CUDA_VERSION))
+    if torch.version.cuda is not None:  # on ROCm we don't want this check
+        CUDA_VERSION = torch._C._cuda_getCompiledVersion()
+        for d in range(device_count()):
+            capability = get_device_capability(d)
+            major = capability[0]
+            minor = capability[1]
+            name = get_device_name(d)
+            if capability == (3, 0) or major < 3:
+                warnings.warn(old_gpu_warn % (d, name, major, capability[1]))
+            elif CUDA_VERSION <= 9000 and major >= 7 and minor >= 5:
+                warnings.warn(incorrect_binary_warn % (d, name, 10000, CUDA_VERSION))
 
 
 def is_initialized():