Fix dvae bug

2021-08-06 14:17:01 -06:00 · 2021-08-06 14:17:01 -06:00 · 0237e96b34
commit 0237e96b34
parent 0799d95af5
1 changed files with 5 additions and 2 deletions
--- 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