Workaround the export issue in torch 2.4

tests/utils/test_cache_utils.py

@@ -174,6 +174,8 @@ def test_static_cache_exportability(self):
         """
         Tests that static cache works with `torch.export()`
         """
+        import torch
+
         if version.parse(torch.__version__) < version.parse("2.3"):
             self.skipTest(reason="This test requires torch >= 2.3 to run.")
 
@@ -217,10 +219,15 @@ def forward(self, tokens: torch.Tensor, input_pos: torch.Tensor):
 
         set_seed(0)
         with torch.no_grad():
-            from torch.export import ExportedProgram, export
+            import torch.export._trace
+            from torch.export import ExportedProgram
 
             model = ExportatibleModelWithStaticCache(config, m)
-            exported_program = export(model, args=(inputs,), kwargs={"input_pos": torch.arange(1)})
+            # Due to issue https://github.com/pytorch/pytorch/issues/128394, we need to switch to use an internal
+            # export API and pre_dispatch=False. Switch to use the public API once the issue is included in 2.4.1+ release.
+            exported_program = torch.export._trace._export(
+                model, args=(inputs,), kwargs={"input_pos": torch.arange(1)}, pre_dispatch=False, strict=True
+            )
             self.assertTrue(isinstance(exported_program, ExportedProgram))