Skip to content

GitLab

Unverified Commit 4e8929dc authored by Wang, Yi's avatar Wang, Yi Committed by GitHub
Browse files

set model to training mode before accelerate.prepare (#24520)

parent 06910f5a
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
src/transformers/trainer.py
View file @ 4e8929dc
...@@ -1640,6 +1640,7 @@ class Trainer: ...@@ -1640,6 +1640,7 @@ class Trainer:
# prepare using `accelerator` prepare # prepare using `accelerator` prepare
if use_accelerator_prepare: if use_accelerator_prepare:
self.model.train()
if hasattr(self.lr_scheduler, "step"): if hasattr(self.lr_scheduler, "step"):
if self.use_apex: if self.use_apex:
model = self.accelerator.prepare(self.model) model = self.accelerator.prepare(self.model)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment