Fix dvae codes issue

This commit is contained in:
James Betker 2021-12-02 23:28:36 -07:00
parent 31fc693a8a
commit 101a01f744

View File

@ -176,7 +176,7 @@ class DiscreteVAE(nn.Module):
return images return images
def get_debug_values(self, step, __): def get_debug_values(self, step, __):
if self.record_codes: if self.record_codes and self.total_codes > 0:
# Report annealing schedule # Report annealing schedule
return {'histogram_codes': self.codes[:self.total_codes]} return {'histogram_codes': self.codes[:self.total_codes]}
else: else: