fix: Make tests pass again

Signed-off-by: Nikos Livathinos <[email protected]>
DS4SD · Jan 27, 2025 · 43ead93 · 43ead93
1 parent 14e71a7
commit 43ead93
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/docling_ibm_models/code_formula_model/models/sam_opt.py b/docling_ibm_models/code_formula_model/models/sam_opt.py
@@ -116,13 +116,13 @@ def forward(
 
                 new_input_embeds.append(cur_input_embeds)
 
-            next_inputs_embeds = torch.stack(new_input_embeds, dim=0)
+            inputs_embeds = torch.stack(new_input_embeds, dim=0)  # type: ignore
 
         return super(SamOPTModel, self).forward(
             input_ids=None,
             attention_mask=attention_mask,
             past_key_values=past_key_values,
-            inputs_embeds=next_inputs_embeds,
+            inputs_embeds=inputs_embeds,
             use_cache=use_cache,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,