fix

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '0.4.4',
+  version = '0.4.5',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   author = 'Phil Wang',

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -335,11 +335,9 @@ def forward(self, x):
                 if self.orthogonal_reg_active_codes_only:
                     # only calculate orthogonal loss for the activated codes for this batch
                     unique_code_ids = torch.unique(embed_ind)
-                    codebook = self.codebook[unique_code_ids]
-                else:
-                    codebook = self.codebook
+                    codebook = codebook[unique_code_ids]
 
-                orthogonal_reg_loss = orthgonal_loss_fn(self.codebook)
+                orthogonal_reg_loss = orthgonal_loss_fn(codebook)
                 loss = loss + orthogonal_reg_loss * self.orthogonal_reg_weight
 
         quantize = self.project_out(quantize)