From 0237e96b34a785d684b3da218197852f1f05eb8e Mon Sep 17 00:00:00 2001 From: James Betker 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