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

…heckpoint
OpenNMT · Nov 19, 2023 · 3c4b7b7 · 3c4b7b7
1 parent fd2f145
commit 3c4b7b7
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)