Problems when running Dola-static decoding

[Lotusboat]

I'm following your work and trying to run your code. However, when I try to run Dola-static mode using the sample "python gsm8k_eval.py --model-name huggyllama/llama-7b --early-exit-layers 16,32 --data-path /path/to/data/folder --output-path output-path.json --num-gpus 1", some errors occured:

Traceback (most recent call last):
File "/mnt/dola/gsm8k_eval.py", line 313, in
model_completion, c_dist = llm.generate(input_text, **generate_kwargs)
File "/mnt/dola/dola.py", line 104, in generate
outputs = self.model.generate(input_ids, max_length=max_len, num_return_sequences=1,
File "/home/miniconda3/envs/dola/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/mnt/dola/transformers-4.28.1/src/transformers/generation/utils.py", line 1244, in generate
self._validate_model_kwargs(model_kwargs.copy())
File "/mnt/dola/transformers-4.28.1/src/transformers/generation/utils.py", line 1115, in _validate_model_kwargs
raise ValueError(
ValueError: The following model_kwargs are not used by the model: ['premature_layer'] (note: typos in the generate arguments will also show up in this list)

My environment has torch=2.7.1 and python=3.10.0. I checked function generate in utils.py but didn't find parameter premature_layers. How can I deal with this? Can you give me some suggestion? Thanks very much!