add grad scaler scale to metrics

2022-03-08 15:52:42 -07:00 · 2022-03-08 15:52:42 -07:00 · 3e5da71b16
commit 3e5da71b16
parent d2bdeb6f20
1 changed files with 3 additions and 1 deletions
--- 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