diff --git a/codes/trainer/steps.py b/codes/trainer/steps.py index c981b87c..b092d093 100644 --- a/codes/trainer/steps.py +++ b/codes/trainer/steps.py @@ -344,4 +344,6 @@ class ConfigurableStep(Module): opt.zero_grad() def get_metrics(self): - return self.loss_accumulator.as_dict() + metrics = self.loss_accumulator.as_dict() + metrics['grad_scaler_scale'] = self.scaler.get_scale() + return metrics