fixup

LoicGrobol · Oct 6, 2023 · 98bcdfb · 98bcdfb
1 parent 006d605
commit 98bcdfb
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/zeldarose/train_transformer.py b/zeldarose/train_transformer.py
@@ -150,9 +150,7 @@ def on_before_zero_grad(
         self,
         trainer: pl.Trainer,
         pl_module: TrainingModule,
-        outputs: pl_types.STEP_OUTPUT,
-        batch: Any,
-        batch_idx: int,
+        optimizer: torch.optim.Optimizer,
     ):
         step = trainer.global_step
         if self.step_period is not None and step % self.step_period == 0: