Support overridden __call__ on nn modules (#138619)

Summary: X-link: pytorch/pytorch#138619 Approved by: https://github.com/williamwen42 ghstack dependencies: #138657 Reviewed By: wdvr Differential Revision: D64881836 Pulled By: anijain2305 fbshipit-source-id: 1974dbc228618e8597eb6ab293272ee985964f52
pytorch · Oct 24, 2024 · b5b342b · b5b342b
1 parent 367b6ef
commit b5b342b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py b/userbenchmark/dynamo/dynamobench/_dynamo/utils.py
@@ -120,6 +120,8 @@
 T = TypeVar("T")
 
 unpatched_nn_module_getattr = torch.nn.Module.__getattr__
+unpatched_nn_module_call = torch.nn.Module.__call__
+unpatched_nn_module_call_impl = torch.nn.Module._call_impl
 
 counters: DefaultDict[str, Counter[str]] = collections.defaultdict(collections.Counter)
 optimus_scuba_log: Dict[str, Any] = {}