From 0237e96b34a785d684b3da218197852f1f05eb8e Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Fri, 6 Aug 2021 14:17:01 -0600
Subject: [PATCH] Fix dvae bug

---
 codes/models/gpt_voice/lucidrains_dvae.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/codes/models/gpt_voice/lucidrains_dvae.py b/codes/models/gpt_voice/lucidrains_dvae.py
index 16e6b148..cd5c25cc 100644
--- a/codes/models/gpt_voice/lucidrains_dvae.py
+++ b/codes/models/gpt_voice/lucidrains_dvae.py
@@ -125,8 +125,11 @@ class DiscreteVAE(nn.Module):
         return images
 
     def get_debug_values(self, step, __):
-        # Report annealing schedule
-        return {'histogram_codes': self.codes}
+        if self.record_codes:
+            # Report annealing schedule
+            return {'histogram_codes': self.codes}
+        else:
+            return {}
 
     @torch.no_grad()
     @eval_decorator