Add --disable-cuda-check option to bypass strict CUDA version validation in setup.py

README.md

@@ -35,6 +35,9 @@ cd apex
 pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
 # otherwise
 pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./
+
+# To prevent cuda checking:
+pip install -v --disable-pip-version-check --disable-cuda-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./
 ```
 
 To reduce the build time of APEX, parallel building can be enhanced via

setup.py

@@ -19,7 +19,8 @@
 
 # ninja build does not work unless include_dirs are abs path
 this_dir = os.path.dirname(os.path.abspath(__file__))
-
+CUDA_DISABLE_CHECK = False
+# Variable to disable cuda torch binary/bare metal version checking with --disable-cuda-check option in pip install
 
 def get_cuda_bare_metal_version(cuda_dir):
     raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
@@ -30,21 +31,22 @@ def get_cuda_bare_metal_version(cuda_dir):
     return raw_output, bare_metal_version
 
 
-def check_cuda_torch_binary_vs_bare_metal(cuda_dir):
+def check_cuda_torch_binary_vs_bare_metal(cuda_dir): 
     raw_output, bare_metal_version = get_cuda_bare_metal_version(cuda_dir)
     torch_binary_version = parse(torch.version.cuda)
 
     print("\nCompiling cuda extensions with")
     print(raw_output + "from " + cuda_dir + "/bin\n")
-
+    if CUDA_DISABLE_CHECK:
+        return
     if (bare_metal_version != torch_binary_version):
         raise RuntimeError(
             "Cuda extensions are being compiled with a version of Cuda that does "
             "not match the version used to compile Pytorch binaries.  "
             "Pytorch binaries were compiled with Cuda {}.\n".format(torch.version.cuda)
             + "In some cases, a minor-version mismatch will not cause later errors:  "
             "https://github.com/NVIDIA/apex/pull/323#discussion_r287021798.  "
-            "You can try commenting out this check (at your own risk)."
+            "You can try commenting out this check (at your own risk). Or use the convenient sys arg --disable-cuda-check with your pip install command!"
         )
 
 
@@ -918,6 +920,8 @@ def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int
     sys.argv.pop(idx + 1)
     sys.argv.pop(idx)
 
+if "--disable-cuda-check" in sys.argv:
+    CUDA_DISABLE_CHECK = True    
 
 # Prevent file conflicts when multiple extensions are compiled simultaneously
 class BuildExtensionSeparateDir(BuildExtension):