diff --git a/beast/train.py b/beast/train.py
index bfa386e..8258d4a 100644
--- a/beast/train.py
+++ b/beast/train.py
@@ -165,7 +165,7 @@ def train(config: dict, model, output_dir: str | Path):
 def get_callbacks(
     checkpointing: bool = True,
     lr_monitor: bool = True,
-    ckpt_every_n_epochs: bool = None,
+    ckpt_every_n_epochs: int | None= None,
 ) -> list:
 
     callbacks = []