Fix #7, remove path to codebook in reconstruction loss

ritheshkumar95 · May 23, 2018 · 8d123c0 · 8d123c0
1 parent cb4702a
commit 8d123c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/modules.py b/modules.py
@@ -80,7 +80,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_, indices = vq_st(z_e_x_, self.embedding.weight)
+        z_q_x_, indices = vq_st(z_e_x_, self.embedding.weight.detach())
         z_q_x = z_q_x_.permute(0, 3, 1, 2).contiguous()
 
         z_q_x_bar_flatten = torch.index_select(self.embedding.weight,