Fix contiguous in VQEmbedding straight-through

Author: tristandeleu

functions.py

@@ -54,7 +54,8 @@ def backward(ctx, grad_output):
             indices, codebook = ctx.saved_tensors
             embedding_size = codebook.size(1)
 
-            grad_output_flatten = grad_output.view(-1, embedding_size)
+            grad_output_flatten = (grad_output.contiguous()
+                                              .view(-1, embedding_size))
             grad_codebook = torch.zeros_like(codebook)
             grad_codebook.index_add_(0, indices, grad_output_flatten)
 

modules.py

@@ -81,7 +81,7 @@ def forward(self, z_e_x):
     def straight_through(self, z_e_x):
         z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()
         z_q_x_ = vq_st(z_e_x_, self.embedding.weight)
-        z_q_x = z_q_x_.permute(0, 3, 1, 2)
+        z_q_x = z_q_x_.permute(0, 3, 1, 2).contiguous()
         return z_q_x