fix(train): use on_fit_start

ZettaAI · Dec 7, 2023 · 9ac6286 · 9ac6286
1 parent 2851b7e
commit 9ac6286
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zetta_utils/training/lightning/trainers/default.py b/zetta_utils/training/lightning/trainers/default.py
@@ -229,7 +229,7 @@ def __init__(
         self.exp_name = exp_name
         self.exp_version = exp_version
 
-    def on_train_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
+    def on_fit_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
         if not os.environ.get("WANDB_MODE", None) == "offline":  # pragma: no cover
             api_key = os.environ.get("WANDB_API_KEY", None)
             wandb.login(key=api_key)