[dynamo] fix potentially missing _torchdynamo_inline from ScriptFunct…

…ion (pytorch#125447) Fix pytorch#119747 Pull Request resolved: pytorch#125447 Approved by: https://github.com/jansel
weiyusheng · May 6, 2024 · bdaa7bb · bdaa7bb
1 parent ad9a27f
commit bdaa7bb
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 0 deletions.
diff --git a/test/dynamo_expected_failures/TestScript.test_nested_breaks b/test/dynamo_expected_failures/TestScript.test_nested_breaks
diff --git a/torch/jit/_script.py b/torch/jit/_script.py
@@ -1391,6 +1391,7 @@ def forward(self, a) -> MyModule:
         _check_directly_compile_overloaded(obj)
         maybe_already_compiled_fn = _try_get_jit_cached_function(obj)
         if maybe_already_compiled_fn:
+            maybe_already_compiled_fn._torchdynamo_inline = obj  # type: ignore[attr-defined]
             return maybe_already_compiled_fn
         ast = get_jit_def(obj, obj.__name__)
         if _rcb is None: