another try

Author: xadupre

_unittests/ut_torch_models/test_llm_phi2.py

@@ -24,12 +24,15 @@ def test_get_phi2(self):
     @requires_torch("2.9.99")
     def test_export_phi2_1_batch_size_1(self):
         # exporting vmap does not work
-        data = get_phi2(num_hidden_layers=2)
+        data = get_phi2(num_hidden_layers=2, batch_size=1)
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        self.assertEqual(inputs["input_ids"].shape[0], 1)
         self.assertEqual(
             {"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
         )
-        with torch_export_patches(patch_transformers=True):
+        with torch.fx.experimental._config.patch(
+            backed_size_oblivious=True
+        ), torch_export_patches(patch_transformers=True):
             ep = torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
             )
@@ -40,8 +43,9 @@ def test_export_phi2_1_batch_size_1(self):
     @requires_torch("2.9.99")
     def test_export_phi2_1_batch_size_2(self):
         # exporting vmap does not work
-        data = get_phi2(num_hidden_layers=2, batch=2)
+        data = get_phi2(num_hidden_layers=2, batch_size=2)
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        self.assertEqual(inputs["input_ids"].shape[0], 2)
         self.assertEqual(
             {"attention_mask", "past_key_values", "input_ids", "position_ids"}, set(inputs)
         )