[dynamo][nn module] Use correct sources for _call_impl (#124970)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/124970
Approved by: https://github.com/jansel
ghstack dependencies: #124779, #124627
This commit is contained in:
Animesh Jain 2024-04-26 11:14:24 -07:00 committed by PyTorch MergeBot
parent 43069c460e
commit fd24d8c05a

View File

@ -340,9 +340,10 @@ class NNModuleVariable(VariableTracker):
# If so at least some changes are needed, we don't allow inlining
# the call_wrapped currently, and maybe other issues too
fn = mod.forward
fn_source = AttrSource(self.source, "forward")
else:
fn = mod._call_impl
fn_source = AttrSource(self.source, "__call__")
fn_source = AttrSource(self.source, "_call_impl")
if istype(fn, types.MethodType):
fn = fn.__func__
fn_source = AttrSource(fn_source, "__func__")