Sum gathered input tokens (#34554)

* sum gathered input tokens * ruff line-length is 119, format the code --------- Co-authored-by: kangsheng <[email protected]>
huggingface · Nov 25, 2024 · 4dc1a69 · 4dc1a69
1 parent 1e492af
commit 4dc1a69
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -2489,7 +2489,9 @@ def _inner_training_loop(
                         else:
                             input_tokens = inputs[main_input_name].numel()
                             input_tokens = torch.tensor(input_tokens, device=self.args.device, dtype=torch.int64)
-                            self.state.num_input_tokens_seen += self.accelerator.gather(input_tokens).cpu().item()
+                            self.state.num_input_tokens_seen += (
+                                self.accelerator.gather(input_tokens).sum().cpu().item()
+                            )
                     if rng_to_sync:
                         self._load_rng_state(resume_from_checkpoint)
                         rng_to_sync = False