diff --git a/examples/callbacks/SAM.py b/examples/callbacks/SAM.py index dbb215b4..93f1c150 100644 --- a/examples/callbacks/SAM.py +++ b/examples/callbacks/SAM.py @@ -36,6 +36,8 @@ }, ) -# run a quick training loop -trainer = pl.Trainer(fast_dev_run=1000, callbacks=[SAM()]) +# run a quick training loop, skipping the first five steps +trainer = pl.Trainer( + fast_dev_run=100, callbacks=[SAM(skip_step_count=5, logging=True, log_level="INFO")] +) trainer.fit(task, datamodule=dm)