Update on "[dynamo][not ready] polyfill infra for classes"

anijain2305 · anijain2305 · commit 7bd428df9ace · 2025-02-09T17:56:29.000-08:00
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames

[ghstack-poisoned]
diff --git a/torch/_dynamo/side_effects.py b/torch/_dynamo/side_effects.py
@@ -297,6 +297,7 @@ def get_variable_cls(self, user_cls):
 
         from .variables.ctx_manager import GenericContextWrappingVariable
         from .variables.torch_function import TorchFunctionModeVariable
+        from .variables.user_defined import is_forbidden_context_manager
 
         variable_cls: type[
             variables.UserDefinedObjectVariable
@@ -305,9 +306,13 @@ def get_variable_cls(self, user_cls):
             user_cls, TorchFunctionMode
         ) and TorchFunctionModeVariable.is_supported_torch_function_mode(user_cls):
             variable_cls = TorchFunctionModeVariable
-        elif hasattr(user_cls, "__enter__") and hasattr(user_cls, "__exit__"):
+        elif (
+            hasattr(user_cls, "__enter__")
+            and hasattr(user_cls, "__exit__")
+            and not is_forbidden_context_manager(user_cls)
+        ):
             variable_cls = GenericContextWrappingVariable
-        if issubclass(user_cls, torch.nn.Module):
+        elif issubclass(user_cls, torch.nn.Module):
             variable_cls = variables.UnspecializedNNModuleVariable
         elif issubclass(user_cls, (dict, collections.OrderedDict)):
             variable_cls = variables.UserDefinedDictVariable
