[Dynamo][Trace PyDispatcher] Remove disable from HigherOrderOperator.__call__ (#146270)

yanboliang · pytorchmergebot · commit bd8d7b1b74b0 · 2025-02-03T21:47:54.000Z
Pull Request resolved: #146270 Approved by: https://github.com/zou3519
diff --git a/test/dynamo/test_higher_order_ops.py b/test/dynamo/test_higher_order_ops.py
@@ -2283,7 +2283,8 @@ def body(x):
 
         res = mod_for_compile(torch.Tensor([[6, 4, 5], [3, 4, 5], [6, 6, 6]]))
         # There is graph break right when we enter body of map
-        self.assertEqual(len(backend.graphs), 0)
+        # Since we are tracing through the Python dispatch logic, it ends up 8 graphs.
+        self.assertEqual(len(backend.graphs), 8)
         self.assertEqual(
             res, mod_for_eager(torch.Tensor([[6, 4, 5], [3, 4, 5], [6, 6, 6]]))
         )
@@ -2319,7 +2320,8 @@ def body(x):
         eager = mod_for_eager(torch.Tensor([[6, 4, 5], [3, 4, 5], [6, 6, 6]]))
         eager = mod_for_eager(torch.Tensor([[6, 4, 5], [3, 4, 5], [6, 6, 6]]))
 
-        self.assertEqual(len(backend.graphs), 0)
+        # Since we are tracing through the Python dispatch logic, it ends up 9 graphs.
+        self.assertEqual(len(backend.graphs), 9)
         self.assertEqual(res, eager)
 
     def test_wrap_subgraph_name_is_valid(self):
diff --git a/torch/_ops.py b/torch/_ops.py
@@ -456,11 +456,6 @@ def check_overloaded(arg):
 
     @abc.abstractmethod
     def __call__(self, /, *args, **kwargs):
-        # Dynamo already traces the body of HigherOrderOp beforehand when it
-        # so no need to trace into it.
-        from torch._dynamo import disable
-
-        @disable
         def wrapper():
             flat_args = _to_flat_tuple(args, kwargs)
             if torch.overrides.has_torch_function(flat_args):
