Issue with "kv_cache" while using modified generate/lora.py for a list of inputs

Dimensions of (k,v) are getting changed due to kv_cache somewhere between line 195-228 in file model.py.

This is happening when line 65 is called by this generate function call

This generate function call I've put under a for loop to iterate over a list of inputs for generation using LoRA weights.

Am trying to find how can I reset this kv_cache, as soon as new input is taken. Any help is appreciated.

More error logs:

Traceback (most recent call last):
  File "generate/lora_itn.py", line 130, in <module>
    CLI(main)
  File "/opt/conda/lib/python3.8/site-packages/jsonargparse/_cli.py", line 85, in CLI
    return _run_component(component, cfg_init)
  File "/opt/conda/lib/python3.8/site-packages/jsonargparse/_cli.py", line 147, in _run_component
    return component(**cfg)
  File "generate/lora_itn.py", line 96, in main
    output = generate(model,
  File "/opt/conda/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/jovyan/mrigank-llm-datavol-1/lit-llama/generate.py", line 71, in generate
    logits = model(x, max_seq_length, input_pos)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/lightning/fabric/wrappers.py", line 116, in forward
    output = self._forward_module(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jovyan/mrigank-llm-datavol-1/lit-llama/lit_llama/model.py", line 114, in forward
    x, self.kv_caches[i] = block(x, rope, mask, max_seq_length, input_pos, self.kv_caches[i])
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jovyan/mrigank-llm-datavol-1/lit-llama/lit_llama/model.py", line 163, in forward
    h, new_kv_cache = self.attn(self.rms_1(x), rope, mask, max_seq_length, input_pos, kv_cache)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/jovyan/mrigank-llm-datavol-1/lit-llama/lit_llama/model.py", line 230, in forward
    y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
RuntimeError: The size of tensor a (318) must match the size of tensor b (314) at non-singleton dimension 3

file lora_itn.py is a modified lora.py with changes for taking in a list of inputs instead.

Start with the kv_cache handling in lit_llama/model.py around lines 195-228 and the cache initialization near line 66. Compare how generate.py line 65 is called from generate/lora.py line 90, then inspect the modified generate/lora_itn.py loop and its traceback. Done means repeated generation over the list of inputs no longer produces the tensor-dimension mismatch.