diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 0bbba3644..30fd8314e 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -102,14 +102,15 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            optimizer_args = {"closure": closure} if closure is not None else {}
+            xm.optimizer_step(self.optimizer, optimizer_args=optimizer_args)
         elif self.scaler is not None:
-            self.scaler.step(self.optimizer)
+            self.scaler.step(self.optimizer, closure)
             self.scaler.update()
         else:
-            self.optimizer.step()
+            self.optimizer.step(closure)
 
     def _switch_parameters(self, parameters_map):
         for param_group in self.optimizer.param_groups:


Review comments: body: ```suggestion
            xm.optimizer_step(self.optimizer, optimizer_args={"closure": closure})
```
The [PyTorch XLA doc](https://pytorch.org/xla/release/1.8/index.html#torch_xla.core.xla_model.optimizer_step) seems to suggest we have to use `optimizer_args` to pass this along and it's not the second argument.
diff_hunk: @@ -102,14 +102,14 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            xm.optimizer_step(self.optimizer, dict(closure=closure))
body: good catch!
diff_hunk: @@ -102,14 +102,14 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            xm.optimizer_step(self.optimizer, dict(closure=closure))
body: Actually thinking a little bit more about it, let's not pass anything if there is no closure:
```suggestion
            optimizer_args = {"closure": closure} if closure is not None else {}
            xm.optimizer_step(self.optimizer, optimizer_args=optimizer_args)
```
diff_hunk: @@ -102,14 +102,14 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            xm.optimizer_step(self.optimizer, optimizer_args={"closure": closure})
body: I concur, this is a little more explicit, but it shouldn't matter (famous last words...) because the base `torch.optim.Optimizer` has a `closure` argument. For the other options, we pass that along anyway. 
diff_hunk: @@ -102,14 +102,14 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            xm.optimizer_step(self.optimizer, optimizer_args={"closure": closure})
body: If you don't mind committing the suggestion then, we're all good to merge :-)
diff_hunk: @@ -102,14 +102,14 @@ def zero_grad(self, set_to_none=None):
                 set_to_none = False
             self.optimizer.zero_grad(set_to_none=set_to_none)
 
-    def step(self):
+    def step(self, closure=None):
         if self.state.distributed_type == DistributedType.TPU:
-            xm.optimizer_step(self.optimizer)
+            xm.optimizer_step(self.optimizer, optimizer_args={"closure": closure})