OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 37 Packages Projects Releases Wiki Activity

[dynamo][inline-inbuilt-nn-modules] Bug fix - Only unspecialized nn modules (#126303)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/126303
Approved by: https://github.com/mlazos, https://github.com/laithsakka
This commit is contained in:
Animesh Jain 2024-05-15 10:55:25 -07:00 committed by PyTorch MergeBot
parent 5cab7a7662
commit bcdd0b11ca
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
torch/_dynamo/mutation_guard.py
View File

@ -103,7 +103,11 @@ def is_dynamic_nn_module(obj, is_export):
# 1) Input signature problem because params are lifted as inputs
# 2) nn module stack info changes
# 3) adjust failing tests
if config.inline_inbuilt_nn_modules and not is_export:
if (
isinstance(obj, torch.nn.Module)
and config.inline_inbuilt_nn_modules
and not is_export
):
return True
dyn = GenerationTracker.dynamic_classes.get(type(obj)) or GenerationTracker.check(
obj