rebase

Author: fegin

torchtitan/distributed/activation_checkpoint.py

@@ -91,13 +91,13 @@ def _apply_op_sac(
 
     def _get_custom_policy(meta):
         def _custom_policy(ctx, func, *args, **kwargs):
-			if (
-				func == torch.ops.aten._to_copy.default
-				and "cuda" in str(args[0].device)
-				and "device" in kwargs
-				and str(kwargs["device"]) == "cpu"
-			):
-				return CheckpointPolicy.MUST_SAVE
+            if (
+                func == torch.ops.aten._to_copy.default
+                and "cuda" in str(args[0].device)
+                and "device" in kwargs
+                and str(kwargs["device"]) == "cpu"
+            ):
+                return CheckpointPolicy.MUST_SAVE
 
             mode = "recompute" if ctx.is_recompute else "forward"
             mm_count_key = f"{mode}_mm_count"