Fixed graph count test because of map autograd rework

pytorch · bohnstingl · May 10, 2025 · May 11, 2025 · May 11, 2025 · May 11, 2025
commit 6bc21e456c6cf40b6fb5157711da93bf187ec582
diff --git a/test/dynamo/test_higher_order_ops.py b/test/dynamo/test_higher_order_ops.py
@@ -2283,8 +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
-        # Since we are tracing through the Python dispatch logic, it ends up 8 graphs.
-        self.assertEqual(len(backend.graphs), 8)
+        # Since we are tracing through the Python dispatch logic, it ends up 3 graphs.
+        self.assertEqual(len(backend.graphs), 3)
         self.assertEqual(
             res, mod_for_eager(torch.Tensor([[6, 4, 5], [3, 4, 5], [6, 6, 6]]))
         )
@@ -2320,8 +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]]))
 
-        # Since we are tracing through the Python dispatch logic, it ends up 9 graphs.
-        self.assertEqual(len(backend.graphs), 9)
+        # Since we are tracing through the Python dispatch logic, it ends up 8 graphs.
+        self.assertEqual(len(backend.graphs), 8)
         self.assertEqual(res, eager)
 
     def test_wrap_subgraph_name_is_valid(self):