Skip to content

Commit 0eb978b

Browse files
committed
another fix
1 parent 31e6d42 commit 0eb978b

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

_unittests/ut_torch_models/test_llm_phi2.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from onnx_diagnostic.torch_models.llms import get_phi2
1010
from onnx_diagnostic.helpers import string_type
1111
from onnx_diagnostic.torch_export_patches import torch_export_patches
12+
from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
1213

1314

1415
class TestLlmPhi(ExtTestCase):
@@ -29,7 +30,9 @@ def test_export_phi2_1(self):
2930
{"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
3031
)
3132
with torch_export_patches(patch_transformers=True):
32-
ep = torch.export.export(model, (), kwargs=inputs, dynamic_shapes=ds)
33+
ep = torch.export.export(
34+
model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
35+
)
3336
assert ep
3437

3538

0 commit comments

Comments
 (0)