initialize the default rank set on TrainerState (#16530)

* initialize the default rank set on TrainerState * fix style

initialize the default rank set on TrainerState (#16530)
* initialize the default rank set on TrainerState * fix style
1cd2e21d · Andres Codas · GitHub · 6f9d8dc1 · 1cd2e21d
Unverified Commit 1cd2e21d authored Apr 04, 2022 by Andres Codas Committed by GitHub Apr 04, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

src/transformers/trainer.py src/transformers/trainer.py +5 -1

No files found.
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -488,7 +488,11 @@ class Trainer:
        else:
            self.label_smoother = None

-        self.state = TrainerState()
+        self.state = TrainerState(
+            is_local_process_zero=self.is_local_process_zero(),
+            is_world_process_zero=self.is_world_process_zero(),
+        )
+
        self.control = TrainerControl()
        # Internal variable to count flos in each process, will be accumulated in `self.state.total_flos` then
        # returned to 0 every time flos need to be logged