fix orthogonal regularization, addressing #43

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.1.5',
+  version = '1.1.6',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -592,15 +592,18 @@ def forward(
             if self.orthogonal_reg_weight > 0:
                 codebook = self._codebook.embed
 
+                # only calculate orthogonal loss for the activated codes for this batch
+
                 if self.orthogonal_reg_active_codes_only:
-                    # only calculate orthogonal loss for the activated codes for this batch
+                    assert not (is_multiheaded and self.separate_codebook_per_head), 'orthogonal regularization for only active codes not compatible with multi-headed with separate codebooks yet'
                     unique_code_ids = torch.unique(embed_ind)
-                    codebook = codebook[unique_code_ids]
+                    codebook = codebook[:, unique_code_ids]
+
+                num_codes = codebook.shape[-2]
 
-                num_codes = codebook.shape[0]
                 if exists(self.orthogonal_reg_max_codes) and num_codes > self.orthogonal_reg_max_codes:
                     rand_ids = torch.randperm(num_codes, device = device)[:self.orthogonal_reg_max_codes]
-                    codebook = codebook[rand_ids]
+                    codebook = codebook[:, rand_ids]
 
                 orthogonal_reg_loss = orthogonal_loss_fn(codebook)
                 loss = loss + orthogonal_reg_loss * self.orthogonal_reg_weight