diff --git a/test/test_cuda.py b/test/test_cuda.py index 62e5f85bb64bca..562558ba75b456 100644 --- a/test/test_cuda.py +++ b/test/test_cuda.py @@ -76,6 +76,7 @@ TestCase, ) from torch.utils.checkpoint import checkpoint_sequential +from torch.utils.cpp_extension import ROCM_HOME from torch.utils.viz._cycles import observe_tensor_cycles @@ -4375,6 +4376,7 @@ def test_mempool_id(self): # increments the id self.assertTrue(abs(pool2[1] - pool1[1]) > 0) + @unittest.skipIf(TEST_WITH_ROCM and not ROCM_HOME, "ROCm requires hipcc compiler") def test_mempool_with_allocator(self): pool = torch.cuda.MemPool()