fix amp wrong call

This commit is contained in:
William Falcon 2020-06-29 06:46:19 -04:00
parent 3ff695510e
commit 593837e1da
1 changed files with 2 additions and 1 deletions

View File

@ -743,7 +743,8 @@ class TrainerTrainLoopMixin(ABC):
# when using 16-bit
else:
native_amp = self.use_amp and NATIVE_AMP_AVALAIBLE
model.optimizer_step(self.current_epoch, batch_idx, optimizer, opt_idx, lambda_closure, native_amp)
model.optimizer_step(self.current_epoch, batch_idx, optimizer, opt_idx, lambda_closure,
using_native_amp=native_amp)
# in native 16-bit we need to update scaler after optimizer step
if self.use_amp and NATIVE_AMP_AVALAIBLE: