OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[dynamo][executorch] Return already added nn.Module during registration (#165338)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165338
Approved by: https://github.com/tugsbayasgalan
This commit is contained in:
Animesh Jain 2025-10-13 10:57:19 -07:00 committed by PyTorch MergeBot
parent ecb53078fa
commit a701c937bf
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
torch/_dynamo/output_graph.py
View File

@ -1079,6 +1079,14 @@ class OutputGraph(OutputGraphCommon):
def register_static_attr_and_return_proxy(
self, attr_prefix: str, attr_value: Any
) -> fx.Proxy:
# Check if the module already exists, if it does, return the already
# added proxy. This is important for executorch tests.
if isinstance(attr_value, torch.nn.Module):
for name, mod in self.nn_modules.items():
if mod is attr_value:
proxy = self.create_proxy("get_attr", name, (), {})
return proxy
attr_name = get_unique_name_wrt(attr_prefix, self.nn_modules)
# TODO `nn_modules` has been historically overloaded to store a lot more
# than just nn module objects, fix that.