skip if no cuda

Author: jerryzh168

test/dtypes/test_affine_quantized.py

@@ -28,7 +28,7 @@
     quantize_,
 )
 from torchao.quantization.quant_primitives import MappingType, ZeroPointDomain
-from torchao.testing.utils import skip_if_rocm
+from torchao.testing.utils import skip_if_no_cuda, skip_if_rocm
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_5,
     TORCH_VERSION_AT_LEAST_2_6,
@@ -308,7 +308,8 @@ def test_alias(self, device, dtype):
         quantize_(dummy, Int8DynamicActivationInt8WeightConfig())
         _ = dummy.weight[...]
 
-    @common_utils.parametrize("device", ["cuda"] if torch.cuda.is_available() else [])
+    @skip_if_no_cuda
+    @common_utils.parametrize("device", ["cuda"])
     @common_utils.parametrize("dtype", [torch.bfloat16])
     def test_slice(self, device, dtype):
         # in_feature not divisible by 1024

torchao/testing/utils.py

@@ -90,6 +90,20 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
+def skip_if_no_cuda():
+    import unittest
+
+    def decorator(test_func):
+        def wrapper(*args, **kwargs):
+            if not torch.cuda.is_available():
+                raise unittest.SkipTest("No cuda available")
+            return test_func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 # copied from https://github.com/pytorch/pytorch/blob/941d094dd1b507dacf06ddc6ed3485a9537e09b7/test/inductor/test_torchinductor.py#L11389
 def copy_tests(my_cls, other_cls, suffix, test_failures=None, xfail_prop=None):  # noqa: B902
     for name, value in my_cls.__dict__.items():