Corrections

arcee-ai · Oct 3, 2023 · 101bcf4 · 101bcf4
1 parent 0de926c
commit 101bcf4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dalm/models/retriever_only_base_model.py b/dalm/models/retriever_only_base_model.py
@@ -37,8 +37,8 @@ def __init__(
     def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
         if self.is_autoregressive:
             # we take the last hidden state of the model
-            token_embeddings = self.model.sample(
-                input_ids, attention_mask, output_hidden_states=True, return_dict_in_generate=True
+            token_embeddings = self.model(
+                input_ids, attention_mask=attention_mask, output_hidden_states=True, return_dict=True
             ).hidden_states[-1]
         else:
             # First element of model_output contains all token embeddings