Fix normalization of the coarse embedding

alexisthual · May 14, 2024 · f8472e3 · f8472e3
1 parent bb3f44f
commit f8472e3
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/src/fugw/scripts/coarse_to_fine.py b/src/fugw/scripts/coarse_to_fine.py
@@ -379,10 +379,16 @@ def fit(
 
     # Normalize embeddings
     source_coarse_embedding = (
-        source_coarse_embedding / source_coarse_embedding.max()
+        source_coarse_embedding
+        / (source_coarse_embedding @ source_coarse_embedding.T)
+        .norm(dim=1)
+        .max()
     )
     target_coarse_embedding = (
-        target_coarse_embedding / target_coarse_embedding.max()
+        target_coarse_embedding
+        / (target_coarse_embedding @ target_coarse_embedding.T)
+        .norm(dim=1)
+        .max()
     )
 
     # Sampled weights