lint

Author: xadupre

_unittests/ut_export/test_serialization.py

@@ -1,6 +1,6 @@
 import unittest
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.ext_test_case import ExtTestCase, requires_transformers
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.helpers.cache_helper import (
     make_dynamic_cache,
@@ -19,6 +19,7 @@ def _get_cache(self, n_layers=2, bsize=2, nheads=4, slen=1, dim=7):
             ]
         )
 
+    @requires_transformers("4.50")
     def test_dynamic_cache(self):
         class Model(torch.nn.Module):
             def forward(self, cache):
@@ -31,6 +32,7 @@ def forward(self, cache):
         exp = torch.export.export(Model(), (cache,), dynamic_shapes=dynamic_shapes)
         self.assertNotEmpty(exp)
 
+    @requires_transformers("4.50")
     def test_dynamic_cache_flat_unflat(self):
         class Model(torch.nn.Module):
             def forward(self, cache):

_unittests/ut_torch_export_patches/test_patch_inputs.py

@@ -1,7 +1,7 @@
 import unittest
 import torch
 import transformers
-from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, requires_transformers
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.torch_export_patches.patch_inputs import (
     convert_dynamic_axes_into_dynamic_shapes,
@@ -10,6 +10,7 @@
 
 class TestPatchInputs(ExtTestCase):
     @hide_stdout()
+    @requires_transformers("4.50")
     def test_convert_dynamic_axes_into_dynamic_shapes_1(self):
         args = (
             torch.randint(0, 10, size=(2, 8)).to(torch.int64),
@@ -55,6 +56,7 @@ def test_convert_dynamic_axes_into_dynamic_shapes_1(self):
         )
 
     @hide_stdout()
+    @requires_transformers("4.50")
     def test_convert_dynamic_axes_into_dynamic_shapes_2(self):
         args = (
             torch.randint(0, 10, size=(2, 8)).to(torch.int64),