From 85863703544b8efc7ff4027c8e94e0c132c1ce91 Mon Sep 17 00:00:00 2001 From: Jeffrey Huynh Date: Tue, 19 Mar 2024 20:15:33 +0000 Subject: [PATCH] Fix indent for 2 lines in ZeRO1 (shard.grad = grad_shard, index += 1) --- torch_xla/distributed/zero_redundancy_optimizer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torch_xla/distributed/zero_redundancy_optimizer.py b/torch_xla/distributed/zero_redundancy_optimizer.py index 2d345081e74..c46642fa448 100644 --- 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