Fix check for backword_pos (#23075)

c6c66584 · Wing Lian · GitHub · f31a510b · c6c66584
Unverified Commit c6c66584 authored May 02, 2023 by Wing Lian Committed by GitHub May 02, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

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

No files found.
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -458,7 +458,9 @@ class Trainer:
                self.fsdp = ShardingStrategy.NO_SHARD
            self.backward_prefetch = BackwardPrefetch.BACKWARD_PRE
-            if "backward_prefetch" in self.args.fsdp_config and "backward_pos" not in self.backward_prefetch:
+            if "backward_prefetch" in self.args.fsdp_config and "backward_pos" in self.args.fsdp_config.get(
+                "backward_prefetch", []
+            ):
                self.backward_prefetch = BackwardPrefetch.BACKWARD_POST
            self.forward_prefetch = False