diff --git a/tests/quantization/finegrained_fp8/test_fp8.py b/tests/quantization/finegrained_fp8/test_fp8.py index e59c2068cde..69881b4cbbf 100644 --- a/tests/quantization/finegrained_fp8/test_fp8.py +++ b/tests/quantization/finegrained_fp8/test_fp8.py @@ -250,6 +250,10 @@ class FP8QuantizerTest(unittest.TestCase): class FP8LinearTest(unittest.TestCase): device = "cuda" + @unittest.skipIf( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] < 9, + "Skipping FP8LinearTest because it is not supported on GPU with capability < 9.0", + ) def test_linear_preserves_shape(self): """ Test that FP8Linear preserves shape when in_features == out_features. @@ -262,6 +266,10 @@ class FP8LinearTest(unittest.TestCase): x_ = linear(x) self.assertEqual(x_.shape, x.shape) + @unittest.skipIf( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] < 9, + "Skipping FP8LinearTest because it is not supported on GPU with capability < 9.0", + ) def test_linear_with_diff_feature_size_preserves_shape(self): """ Test that FP8Linear generates the correct shape when in_features != out_features.