Fixed multiple save when using step-save-period

in conjunction with bach accumulation (close #30)
LoicGrobol · Oct 6, 2023 · 006d605 · 006d605
1 parent edb2613
commit 006d605
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,11 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/) and this p
 
 ## [Unreleased]
 
+### Fixed
+
+- Fixed multiple save when using step-save-period in conjunction with bach accumulation (close issue
+  [#30](https://github.com/LoicGrobol/zeldarose/issues/30))
+
 ### Changed
 
 - Maximum Pyorch compatibility bumped to 2.1

diff --git a/zeldarose/train_transformer.py b/zeldarose/train_transformer.py
@@ -146,7 +146,7 @@ def __init__(
         self.tokenizer = tokenizer
 
     @rank_zero_only
-    def on_train_batch_end(
+    def on_before_zero_grad(
         self,
         trainer: pl.Trainer,
         pl_module: TrainingModule,