[CP] Fix incorrect indentation

fegin · fegin · commit ea4b7d7208ef · 2025-10-14T23:10:56.000-07:00
Introduced by #1776. Verified with the comment: ``` CONFIG_FILE="./torchtitan/models/llama3/train_configs/debug_model.toml" ./run_train.sh --training.steps=10 --parallelism.context_parallel_degree=8 ```
diff --git a/torchtitan/distributed/utils.py b/torchtitan/distributed/utils.py
@@ -207,6 +207,7 @@ def context(cp_context: Generator[None, None, None] | None = None):
                     torch._dynamo.utils.maybe_enable_compiled_autograd(True)
                 )
 
+            if cp_context:
                 stack.enter_context(cp_context)
 
             yield

Original file line number	Diff line number	Diff line change
`@@ -207,6 +207,7 @@ def context(cp_context: Generator[None, None, None] \| None = None):`
`207`	`207`	`torch._dynamo.utils.maybe_enable_compiled_autograd(True)`
`208`	`208`	`)`
`209`	`209`
	`210`	`+ if cp_context:`
`210`	`211`	`stack.enter_context(cp_context)`
`211`	`212`
`212`	`213`	`yield`