Error when use llama-1-160m

[Chlience]

Error info:

Traceback (most recent call last):
  File "/home/data2/wsq/ouroboros/test.py", line 29, in <module>
    ouroboros_output = ouroboros(input_ids, small_model, target_model, max_len=80, gamma=gamma, window_size=window_size, guess_set_size=guess_set_size, lookahead_level=lookahead_level)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/data2/wsq/miniconda3/envs/ouroboros/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/data2/wsq/ouroboros/ouroboros/ouroboros.py", line 92, in ouroboros
    approx_model_cache.rollback(n+1)
  File "/home/data2/wsq/miniconda3/envs/ouroboros/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/data2/wsq/ouroboros/ouroboros/kv_cache_model.py", line 44, in rollback
    k = self.ctx['past_key_values'][i][0][:,:,:end_pos,:]
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^
IndexError: list index out of range

[huangyuxiang03]

Thanks for fixing this. We will merge the fixed bug as long as we are free.