refactor: adding step configuration to SAM example

Signed-off-by: Lee, Kin Long Kelvin <[email protected]>
IntelLabs · Sep 10, 2024 · ff0a415 · ff0a415
1 parent 60a04e7
commit ff0a415
Showing 1 changed file with 4 additions and 2 deletions.
diff --git 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)