From 3e5da71b16bb092cac0a251d44b7e1556845cd91 Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Tue, 8 Mar 2022 15:52:42 -0700
Subject: [PATCH] add grad scaler scale to metrics

---
 codes/trainer/steps.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

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