Skip to content

Commit c3376bf

Browse files
committed
fix version compitablity
Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
1 parent 2d65a72 commit c3376bf

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

optimum/intel/ipex/modeling_base.py

+4-3
Original file line numberDiff line numberDiff line change
@@ -353,7 +353,7 @@ def _reorder_cache(self, *args, **kwargs):
353353
def prepare_inputs_for_generation(self, *args, **kwargs):
354354
return self.model.prepare_inputs_for_generation(*args, **kwargs)
355355

356-
def _supports_logits_to_keep(self) -> bool:
356+
def _supports_num_logits_to_keep(self) -> bool:
357357
"""
358358
Return True if the current model supports the keyword argument `logits_to_keep` in forward()
359359
to save memory. Checking it in this way allows to avoid using a new model attribute.
@@ -471,12 +471,13 @@ def prepare_inputs_for_generation(self, *args, **kwargs):
471471
def get_encoder(self, *args, **kwargs):
472472
return self.model.get_encoder(*args, **kwargs)
473473

474-
def _supports_logits_to_keep(self) -> bool:
474+
def _supports_num_logits_to_keep(self) -> bool:
475475
"""
476476
Return True if the current model supports the keyword argument `logits_to_keep` in forward()
477477
to save memory. Checking it in this way allows to avoid using a new model attribute.
478478
"""
479-
return "logits_to_keep" in set(inspect.signature(self.model.forward).parameters.keys())
479+
logits_to_keep_name = "logits_to_keep" if is_transformers_version(">", "4.49") else "num_logits_to_keep"
480+
return logits_to_keep_name in set(inspect.signature(self.model.forward).parameters.keys())
480481

481482
def _init_warmup(self):
482483
inputs = prepare_jit_inputs(self.model, self.export_feature, False)

0 commit comments

Comments
 (0)