Fix error to load data at the correct position when resuming from a c…

…heckpoint
OpenNMT · Jan 7, 2024 · 0a06542 · 0a06542
1 parent a987e7d
commit 0a06542
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/onmt/trainer.py b/onmt/trainer.py
@@ -306,8 +306,11 @@ def train(
         torch.cuda.empty_cache()
 
         for i, (batches, normalization) in enumerate(self._accum_batches(train_iter)):
-
             step = self.optim.training_step
+            # Check if the iterator is at the same step as the optimizer
+            if i + 1 != step:
+                continue
+
             # UPDATE DROPOUT
             self._maybe_update_dropout(step)