fix when attention_mask=None

Signed-off-by: jiqing-feng <[email protected]>
huggingface · Dec 12, 2024 · ba910ed · ba910ed
1 parent bb51139
commit ba910ed
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/optimum/intel/ipex/modeling_base.py b/optimum/intel/ipex/modeling_base.py
@@ -276,6 +276,8 @@ def forward(
         attention_mask: Optional[torch.FloatTensor] = None,
         **kwargs,
     ) -> CausalLMOutputWithPast:
+        if self.add_patch and input_ids is not None and attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
         return self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
 
     def _prepare_generation_config(