test: squeeze-21

ENOT-AutoDL · Jul 30, 2024 · 9a58fec · 9a58fec
1 parent 3dabce4
commit 9a58fec
Showing 1 changed file with 13 additions and 22 deletions.
diff --git a/tests/node_converters/squeeze_test.py b/tests/node_converters/squeeze_test.py
@@ -55,30 +55,21 @@ def _test_squeeze(
 
 
 @pytest.mark.filterwarnings('ignore::torch.jit._trace.TracerWarning')
+@pytest.mark.parametrize('opset_version', [11, 13, 21])
 @pytest.mark.parametrize(
-    'shape,axes,opset_version',
+    'shape, axes',
     (
-        ([1, 3, 4, 5], [0], 11),
-        ([1, 3, 1, 5], [-2], 11),
-        ([1, 3, 1, 5], [0, 2], 11),
-        ([1, 3, 1, 5], [2, 0], 11),
-        ([1, 3, 1, 1, 1, 5, 1], [2, 0, 6], 11),
-        ([1, 3, 1, 5], [0, -2], 11),
-        ([1, 3, 1, 5], [-2, 0], 11),
-        ([1, 3, 1, 5], None, 11),
-        ([1, 1, 1, 1], None, 11),
-        ([1], None, 11),
-        ([3, 3, 3], None, 11),
-        ([1, 3, 4, 5], [0], 13),
-        ([1, 3, 1, 5], [-2], 13),
-        ([1, 3, 1, 5], [0, 2], 13),
-        ([1, 3, 1, 5], [2, 0], 13),
-        ([1, 3, 1, 5], [0, -2], 13),
-        ([1, 3, 1, 5], [-2, 0], 13),
-        ([1, 3, 1, 5], None, 13),
-        ([1, 1, 1, 1], None, 13),
-        ([1], None, 13),
-        ([3, 3, 3], None, 13),
+        ([1, 3, 4, 5], [0]),
+        ([1, 3, 1, 5], [-2]),
+        ([1, 3, 1, 5], [0, 2]),
+        ([1, 3, 1, 5], [2, 0]),
+        ([1, 3, 1, 1, 1, 5, 1], [2, 0, 6]),
+        ([1, 3, 1, 5], [0, -2]),
+        ([1, 3, 1, 5], [-2, 0]),
+        ([1, 3, 1, 5], None),
+        ([1, 1, 1, 1], None),
+        ([1], None),
+        ([3, 3, 3], None),
     ),
 )
 def test_squeeze(  # pylint: disable=missing-function-docstring