Forgot the trainer ups

TJ-Solergibert · Jul 16, 2024 · efe8720 · efe8720
1 parent d9f0670
commit efe8720
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/nanotron/trainer.py b/src/nanotron/trainer.py
@@ -393,7 +393,10 @@ def find_stage_idx_to_resume():
 
     def train(
         self,
-        dataloader_or_dls: Dict[
+        train_dataloader_or_dls: Dict[
+            str, Union[Iterator[Dict[str, Union[torch.Tensor, TensorPointer]]], Tuple[Iterator, ...]]
+        ],
+        valid_dataloader_or_dls: Dict[
             str, Union[Iterator[Dict[str, Union[torch.Tensor, TensorPointer]]], Tuple[Iterator, ...]]
         ],
         **kwargs,
@@ -424,7 +427,7 @@ def train(
                     prof.step()
 
                 self.iteration_start_time = time.time()
-                self._update_dataloader_based_on_training_stages(dataloader_or_dls)
+                self._update_dataloader_based_on_training_stages(train_dataloader_or_dls)
 
                 # Training step
                 outputs, loss_avg = self.training_step(dataloader=self.current_dataloader)