fix matryoshka norm

DesmonDay · Jan 13, 2025 · 1cf0578 · 1cf0578
1 parent 7c3d9d5
commit 1cf0578
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/paddlenlp/transformers/contrastive_loss.py b/paddlenlp/transformers/contrastive_loss.py
@@ -52,10 +52,10 @@ def forward(self, q_reps, p_reps):
         if len(self.embedding_matryoshka_dims) > 0:
             loss = 0.0
             for dim in self.embedding_matryoshka_dims:
-                reduced_q_reps = q_reps[:, :dim]
+                reduced_q_reps = q_reps[:, :dim].astype("float32")
                 reduced_q_reps = nn.functional.normalize(reduced_q_reps, axis=-1)
 
-                reduced_p_reps = p_reps[:, :dim]
+                reduced_p_reps = p_reps[:, :dim].astype("float32")
                 reduced_p_reps = nn.functional.normalize(reduced_p_reps, axis=-1)
 
                 dim_loss = self.loss_fn(reduced_q_reps, reduced_p_reps)