Skip to content

Commit

Permalink
Fix callback call in Fabric Trainer example (#19986)
Browse files Browse the repository at this point in the history
  • Loading branch information
liambsmith committed Jun 18, 2024
1 parent c1af4d0 commit 394c42a
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion examples/fabric/build_your_own_trainer/trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,7 +227,7 @@ def train_loop(
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))
Expand Down

0 comments on commit 394c42a

Please sign in to comment.