Fix dvae bug

This commit is contained in:
James Betker 2021-08-06 14:17:01 -06:00
parent 0799d95af5
commit 0237e96b34

View File

@ -125,8 +125,11 @@ class DiscreteVAE(nn.Module):
return images return images
def get_debug_values(self, step, __): def get_debug_values(self, step, __):
# Report annealing schedule if self.record_codes:
return {'histogram_codes': self.codes} # Report annealing schedule
return {'histogram_codes': self.codes}
else:
return {}
@torch.no_grad() @torch.no_grad()
@eval_decorator @eval_decorator