Merge pull request #156 from lweitkamp/master

Remove double l2 normalization
lucidrains · Aug 27, 2024 · 15629a5 · 15629a5
2 parents 1447998 + ab5a61a
commit 15629a5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vector_quantize_pytorch/vector_quantize_pytorch.py b/vector_quantize_pytorch/vector_quantize_pytorch.py
@@ -751,7 +751,7 @@ def forward(
             embed_normalized = self.embed_avg / rearrange(cluster_size, '... -> ... 1')
             embed_normalized = l2norm(embed_normalized)
 
-            self.embed.data.copy_(l2norm(embed_normalized))
+            self.embed.data.copy_(embed_normalized)
             self.expire_codes_(x)
 
         if needs_codebook_dim: