Merge in main

Signed-off-by: Vibhu Jawa <[email protected]>
NVIDIA · Feb 5, 2025 · d876faa · d876faa
1 parent 3cb6354
commit d876faa
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_semdedup.py b/tests/test_semdedup.py
@@ -154,7 +154,7 @@ def get_reference_embeddings(
 
         if pooling_strategy == "last_token":
             embeddings = outputs.last_hidden_state[:, -1, :]
-        elif pooling_strategy == "mean":
+        elif pooling_strategy == "mean_pooling":
             token_embeddings = outputs.last_hidden_state
             attention_mask = inputs["attention_mask"]
             input_mask_expanded = (