Temporary Fix: Skip on H100

pytorch · Oct 3, 2024 · 2b83972 · 2b83972
1 parent 09b8b3c
commit 2b83972
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/test/dtypes/test_affine_quantized_tensor_parallel.py b/test/dtypes/test_affine_quantized_tensor_parallel.py
@@ -5,8 +5,12 @@
 class TestAffineQuantizedTensorParallel(TorchAOTensorParallelTestCase):
     pass
 
+is_H100 = torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0)
 
-copy_tests(TorchAOTensorParallelTestCase, TestAffineQuantizedTensorParallel, "aqt_tp")
+if not is_H100:
+    copy_tests(TorchAOTensorParallelTestCase, TestAffineQuantizedTensorParallel, "aqt_tp")
+else:
+    print("Skipping TestAffineQuantizedTensorParallel because it doesn't run on H100")
 
 if __name__ == "__main__":
     run_tests()