Fix 2.5.1 failing sparsity test (#1261)

wip
pytorch · Nov 11, 2024 · a827d04 · a827d04
1 parent 75f52ae
commit a827d04
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/prototype/test_sparse_api.py b/test/prototype/test_sparse_api.py
@@ -18,6 +18,7 @@
     TORCH_VERSION_AT_LEAST_2_3,
     TORCH_VERSION_AT_LEAST_2_4,
     TORCH_VERSION_AT_LEAST_2_5,
+    TORCH_VERSION_AT_LEAST_2_6,
 )
 
 
@@ -161,7 +162,7 @@ def test_sparse(self, compile):
 
 
 class TestQuantBlockSparseWeight(common_utils.TestCase):
-    @unittest.skipIf(not TORCH_VERSION_AFTER_2_5, "pytorch 2.6+ feature")
+    @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_6, "pytorch 2.6+ feature")
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
     @common_utils.parametrize("compile", [True, False])
     def test_sparse(self, compile):