Fix indent for 2 lines in ZeRO1 (shard.grad = grad_shard, index += 1)

pytorch · Mar 19, 2024 · 8586370 · 8586370
1 parent ae348b2
commit 8586370
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_xla/distributed/zero_redundancy_optimizer.py b/torch_xla/distributed/zero_redundancy_optimizer.py
@@ -314,8 +314,8 @@ def step(self, closure=None, **kwargs):
 
             if grad_shard.dtype != self.optimizer_dtype:
               grad_shard = grad_shard.to(dtype=self.optimizer_dtype)
-          shard.grad = grad_shard
-          index += 1
+            shard.grad = grad_shard
+            index += 1
 
     if self.grad_clipping:
       # Update unscale/clip with sub partitions