Fix quantizer with balancing_heuristic

codes/models/vqvae/vqvae.py

@@ -48,7 +48,7 @@ class Quantize(nn.Module):
         self.register_buffer("embed_avg", embed.clone())
 
     def forward(self, input):
-        if self.codes_full:
+        if self.balancing_heuristic and self.codes_full:
             h = torch.histc(self.codes, bins=self.n_embed, min=0, max=self.n_embed) / len(self.codes)
             mask = torch.logical_or(h > .9, h < .01).unsqueeze(1)
             ep = self.embed.permute(1,0)
@@ -73,6 +73,7 @@ class Quantize(nn.Module):
         embed_ind = embed_ind.view(*input.shape[:-1])
         quantize = self.embed_code(embed_ind)
 
+        if self.balancing_heuristic:
             if self.codes is None:
                 self.codes = embed_ind.flatten()
             else: