fix TEST_ROCM definition to disable test_jit_cudnn_extension on rocm (#110385)

dnikolaev-amd · pytorchmergebot · commit 901aa85b58e8 · 2023-10-04T20:02:02.000Z
Define TEST_ROCM before modification TEST_CUDA. Otherwise TEST_ROCM will always be false and will not disable test_jit_cudnn_extension for rocm Fixes #107182 Pull Request resolved: #110385 Approved by: https://github.com/jithunnair-amd, https://github.com/kit1980
diff --git a/test/test_cpp_extensions_jit.py b/test/test_cpp_extensions_jit.py
@@ -20,8 +20,9 @@
 import torch.multiprocessing as mp
 from torch.utils.cpp_extension import _TORCH_PATH, remove_extension_h_precompiler_headers, get_cxx_compiler, check_compiler_is_gcc
 
-TEST_CUDA = TEST_CUDA and CUDA_HOME is not None
+# define TEST_ROCM before changing TEST_CUDA
 TEST_ROCM = TEST_CUDA and torch.version.hip is not None and ROCM_HOME is not None
+TEST_CUDA = TEST_CUDA and CUDA_HOME is not None
 TEST_MPS = torch.backends.mps.is_available()
 IS_WINDOWS = sys.platform == "win32"
 IS_LINUX = sys.platform.startswith('linux')
