device

Author: jerryzh168

torchao/testing/utils.py

@@ -165,11 +165,9 @@ def test_linear_compile(self, device, dtype):
     NUM_DEVICES,
 )
 
-@unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
 class TorchAOTensorParallelTestCase(DTensorTestBase):
     """Basic test case for tensor subclasses
     """
-    COMMON_DEVICES = (["cuda"] if torch.cuda.is_available() else [])
     COMMON_DTYPES = [torch.float32, torch.float16, torch.bfloat16]
 
     TENSOR_SUBCLASS = AffineQuantizedTensor
@@ -221,10 +219,11 @@ def quantize(self, m: torch.nn.Module) -> torch.nn.Module:
         quantize_(m, self.QUANT_METHOD_FN(**self.QUANT_METHOD_KWARGS))
         return m
 
-    @common_utils.parametrize("device", COMMON_DEVICES)
     @common_utils.parametrize("dtype", COMMON_DTYPES)
     @with_comms
-    def test_tp(self, device, dtype):
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
+    def test_tp(self, dtype):
+        device = "cuda"
         # To make sure different ranks create the same module
         torch.manual_seed(5)