deepspeed + grad acumm (#9622)

This commit is contained in:
Stas Bekman
2021-01-15 10:12:26 -08:00
committed by GitHub
parent 6d3b688b04
commit c60e0e1ee4
2 changed files with 8 additions and 1 deletions

View File

@@ -931,7 +931,9 @@ class Trainer:
)
# Optimizer step
if is_torch_tpu_available():
if self.deepspeed:
self.deepspeed.step()
elif is_torch_tpu_available():
xm.optimizer_step(self.optimizer)
elif self.use_amp:
self.scaler.step(self.optimizer)