@@ -353,12 +353,12 @@ def _reorder_cache(self, *args, **kwargs):
353
353
def prepare_inputs_for_generation (self , * args , ** kwargs ):
354
354
return self .model .prepare_inputs_for_generation (* args , ** kwargs )
355
355
356
- def _supports_num_logits_to_keep (self ) -> bool :
356
+ def _supports_logits_to_keep (self ) -> bool :
357
357
"""
358
- Return True if the current model supports the keyword argument `num_logits_to_keep ` in forward()
358
+ Return True if the current model supports the keyword argument `logits_to_keep ` in forward()
359
359
to save memory. Checking it in this way allows to avoid using a new model attribute.
360
360
"""
361
- return "num_logits_to_keep " in set (inspect .signature (self . model .forward ).parameters .keys ())
361
+ return "logits_to_keep " in set (inspect .signature (self .forward ).parameters .keys ())
362
362
363
363
def generate (self , * args , ** kwargs ):
364
364
if self ._add_patch and kwargs .get ("assistant_model" , None ):
@@ -470,12 +470,12 @@ def prepare_inputs_for_generation(self, *args, **kwargs):
470
470
def get_encoder (self , * args , ** kwargs ):
471
471
return self .model .get_encoder (* args , ** kwargs )
472
472
473
- def _supports_num_logits_to_keep (self ) -> bool :
473
+ def _supports_logits_to_keep (self ) -> bool :
474
474
"""
475
- Return True if the current model supports the keyword argument `num_logits_to_keep ` in forward()
475
+ Return True if the current model supports the keyword argument `logits_to_keep ` in forward()
476
476
to save memory. Checking it in this way allows to avoid using a new model attribute.
477
477
"""
478
- return "num_logits_to_keep " in set (inspect .signature (self . model .forward ).parameters .keys ())
478
+ return "logits_to_keep " in set (inspect .signature (self .forward ).parameters .keys ())
479
479
480
480
def _init_warmup (self ):
481
481
inputs = prepare_jit_inputs (self .model , self .export_feature , False )
0 commit comments