From ee2827dee9134f5a1072da9022568952e5e9834a Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Fri, 10 Jun 2022 15:23:31 -0600
Subject: [PATCH] Debug warmup state

---
 codes/trainer/injectors/gaussian_diffusion_injector.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/codes/trainer/injectors/gaussian_diffusion_injector.py b/codes/trainer/injectors/gaussian_diffusion_injector.py
index 0f1e1a54..7ddceeb3 100644
--- a/codes/trainer/injectors/gaussian_diffusion_injector.py
+++ b/codes/trainer/injectors/gaussian_diffusion_injector.py
@@ -56,6 +56,13 @@ class GaussianDiffusionInjector(Injector):
             self.channel_balancing_fn = None
         assert k <= 1, 'Only one channel filtering function can be applied.'
 
+    def extra_metrics(self):
+        if hasattr(self, 'schedule_sampler') and isinstance(self.schedule_sampler, LossSecondMomentResampler):
+            return {
+                'sampler_warmed_up': self.schedule_sampler._warmed_up()
+            }
+        return {}
+
     def forward(self, state):
         gen = self.env['generators'][self.opt['generator']]
         hq = state[self.input]