bring back the other half of the commit loss even in the presence of rotation trick, addressing #177

Author: lucidrains

examples/autoencoder_sim_vq.py

@@ -16,7 +16,7 @@
 seed = 1234
 
 rotation_trick = True  # rotation trick instead ot straight-through
-use_mlp = True        # use a one layer mlp with relu instead of linear
+use_mlp = True         # use a one layer mlp with relu instead of linear
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 

vector_quantize_pytorch/sim_vq.py

@@ -118,18 +118,15 @@ def forward(
 
         # commit loss and straight through, as was done in the paper
 
-        commit_loss = F.mse_loss(x.detach(), quantized)
+        commit_loss = (
+            F.mse_loss(x.detach(), quantized) +
+            F.mse_loss(x, quantized.detach()) * self.input_to_quantize_commit_loss_weight
+        )
 
         if self.rotation_trick:
             # rotation trick from @cfifty
             quantized = rotate_to(x, quantized)
         else:
-
-            commit_loss = (
-                commit_loss + 
-                F.mse_loss(x, quantized.detach()) * self.input_to_quantize_commit_loss_weight
-            )
-
             quantized = (quantized - x).detach() + x
 
         quantized = inverse_pack(quantized)