From 0164560353aab22b9c7769d344f833ae8a16f478 Mon Sep 17 00:00:00 2001 From: Guang Yang <42389959+guangy10@users.noreply.github.com> Date: Thu, 8 Aug 2024 10:13:40 -0700 Subject: [PATCH] Fixed test `test_static_cache_exportability` with torch 2.4.0 (#32516) Workaround the export issue in torch 2.4 Co-authored-by: Guang Yang --- tests/utils/test_cache_utils.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py index e8d120216a..5b61999b56 100644 --- a/tests/utils/test_cache_utils.py +++ b/tests/utils/test_cache_utils.py @@ -174,6 +174,8 @@ class CacheTest(unittest.TestCase): """ 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 @@ class CacheTest(unittest.TestCase): 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))