fix(sft_trainer): total_steps calculation when distributed (#432)

CarperAI · Apr 13, 2023 · 92b68e4 · 92b68e4
1 parent adb3be2
commit 92b68e4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trlx/trainer/accelerate_sft_trainer.py b/trlx/trainer/accelerate_sft_trainer.py
@@ -64,7 +64,7 @@ def prepare_learning(self):
         ) = self.accelerator.prepare(self.model, self.opt, train_dataloader, eval_dataloader)
 
         self.n_updates_per_batch = 1
-        self.total_steps = self.config.train.epochs * len(train_dataloader)
+        self.total_steps = self.config.train.epochs * len(self.train_dataloader)
         self.total_steps = min(self.total_steps, self.config.train.total_steps)
 
     def make_experience(self, samples, seq_length):