Fix check for backword_pos (huggingface#23075)

hxfxjun · May 2, 2023 · c6c6658 · c6c6658
1 parent f31a510
commit c6c6658
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -458,7 +458,9 @@ def __init__(
                 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