Don't pass tuple to with statement (pytorch#110864)

Signed-off-by: Edward Z. Yang <[email protected]> Pull Request resolved: pytorch#110864 Approved by: https://github.com/Skylion007, https://github.com/awgu
kurtamohler · Oct 9, 2023 · 8ae623d · 8ae623d
1 parent 4b881b0
commit 8ae623d
Showing 1 changed file with 8 additions and 14 deletions.
diff --git a/test/distributed/test_inductor_collectives.py b/test/distributed/test_inductor_collectives.py
@@ -314,13 +314,10 @@ def example(inp, input_split_sizes_tensor, output_split_sizes_tensor, *, tag, ra
             out = a2a / a2a.sum(dim=0)
             return out
 
-        with (
-            _dynamo_dist_per_rank_init(self.rank, self.world_size),
-            torch._dynamo.config.patch(
-                dynamic_shapes=True,
-                capture_dynamic_output_shape_ops=True,
-                capture_scalar_outputs=True,
-            )
+        with _dynamo_dist_per_rank_init(self.rank, self.world_size), torch._dynamo.config.patch(
+            dynamic_shapes=True,
+            capture_dynamic_output_shape_ops=True,
+            capture_scalar_outputs=True,
         ):
             row = self.world_size * (self.rank + 1) * (self.world_size + 1) / 2
             input_split_sizes_tensor = torch.tensor([(i + 1) * (self.rank + 1) for i in range(self.world_size)], dtype=torch.int64)
@@ -400,13 +397,10 @@ def example(inp, output_split_sizes_tensor, *, tag, ranks, group_size):
             out = a2a / a2a.sum(dim=0)
             return out
 
-        with (
-            _dynamo_dist_per_rank_init(self.rank, self.world_size),
-            torch._dynamo.config.patch(
-                dynamic_shapes=True,
-                capture_dynamic_output_shape_ops=True,
-                capture_scalar_outputs=True,
-            )
+        with _dynamo_dist_per_rank_init(self.rank, self.world_size), torch._dynamo.config.patch(
+            dynamic_shapes=True,
+            capture_dynamic_output_shape_ops=True,
+            capture_scalar_outputs=True,
         ):
             output_split_sizes_tensor = torch.tensor([1] * self.world_size, dtype=torch.int64)
             inputs = (