Updated test

MinishLab · Oct 11, 2024 · 27d0a39 · 27d0a39
1 parent 826f01f
commit 27d0a39
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/model2vec/distill/inference.py b/model2vec/distill/inference.py
@@ -144,9 +144,10 @@ def create_output_embeddings_from_model_name(
             if out.dtype == torch.bfloat16:
                 out = out.float()
 
-            # Detach the tensor from the computation graph before converting to NumPy
+        # Add the output to the intermediate weights
         intermediate_weights.append(out[:, 1].detach().cpu().numpy())
 
+    # Concatenate the intermediate weights
     out_weights = np.concatenate(intermediate_weights)
 
     return tokenizer.convert_ids_to_tokens(ids), out_weights