From 394c42aaf6f5f67ad4cde338428c2b1cdd03c882 Mon Sep 17 00:00:00 2001 From: liambsmith Date: Tue, 18 Jun 2024 13:14:32 -0400 Subject: [PATCH] Fix callback call in Fabric Trainer example (#19986) --- examples/fabric/build_your_own_trainer/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/fabric/build_your_own_trainer/trainer.py b/examples/fabric/build_your_own_trainer/trainer.py index c5e6836bfe..a225bf5556 100644 --- a/examples/fabric/build_your_own_trainer/trainer.py +++ b/examples/fabric/build_your_own_trainer/trainer.py @@ -227,7 +227,7 @@ class MyCustomTrainer: should_optim_step = self.global_step % self.grad_accum_steps == 0 if should_optim_step: # currently only supports a single optimizer - self.fabric.call("on_before_optimizer_step", optimizer, 0) + self.fabric.call("on_before_optimizer_step", optimizer) # optimizer step runs train step internally through closure optimizer.step(partial(self.training_step, model=model, batch=batch, batch_idx=batch_idx))