fsdp fix (#20719)

a413c725 · Sourab Mangrulkar · GitHub · 17c742bb · a413c725
Unverified Commit a413c725 authored Dec 12, 2022 by Sourab Mangrulkar Committed by GitHub Dec 12, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 2 deletions

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

No files found.
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1424,9 +1424,8 @@ class Trainer:
                    cpu_offload=cpu_offload,
                    auto_wrap_policy=auto_wrap_policy,
                    mixed_precision=mixed_precision_policy,
+                    device_id=self.args.device,
                )
-                if FSDPOption.OFFLOAD not in self.args.fsdp:
-                    model.to(self.args.device)
        elif is_sagemaker_dp_enabled():
            model = nn.parallel.DistributedDataParallel(
                model, device_ids=[int(os.getenv("SMDATAPARALLEL_LOCAL_RANK"))]