Add kwargs back to the closure (mosaicml#292)

* add kwargs back to the closure * woops
coryMosaicML · Feb 23, 2022 · a6bb3fd · a6bb3fd
1 parent 196cfc4
commit a6bb3fd
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
@@ -636,10 +636,11 @@ def _train_loop(self) -> None:
                     elif self._use_closures():
                         for optimizer in state.optimizers:
                             if use_grad_scaling:
-                                total_loss = state.scaler.step(optimizer,
-                                                               closure=lambda: self._train_batch(microbatches))
+                                total_loss = state.scaler.step(
+                                    optimizer, closure=lambda **kwargs: self._train_batch(microbatches, **kwargs))
                             else:
-                                total_loss = optimizer.step(closure=lambda: self._train_batch(microbatches).item())
+                                total_loss = optimizer.step(
+                                    closure=lambda **kwargs: self._train_batch(microbatches, **kwargs).item())
                     else:
                         total_loss = self._train_batch(microbatches)
                         for optimizer in state.optimizers: