Working with 8bit and 4bit quantized models

Hey! Great work on this project! I got it t work on a couple of t5 instruction tuned models from huggingface, I was just curious, has anyone been able to get the code to work with quantized modes? Currently when I set it to 'load_in_4bit=True' I get this error:

`╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ in :1 │
│ │
│ /content/unlimiformer/src/unlimiformer.py:707 in convert_model │
│ │
│ 704 │ @classmethod │
│ 705 │ def convert_model(cls, model, *args, **kwargs): │
│ 706 │ │ model_clone = AutoModelForSeq2SeqLM.from_config(model.config) │
│ ❱ 707 │ │ model_clone.load_state_dict(model.state_dict()) │
│ 708 │ │ type_to_class = { │
│ 709 │ │ │ BartModel: UnlimiformerBART, │
│ 710 │ │ │ BartForConditionalGeneration: UnlimiformerBART, │
│ │
│ /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:2041 in load_state_dict │
│ │
│ 2038 │ │ │ │ │ │ ', '.join('"{}"'.format(k) for k in missing_keys))) │
│ 2039 │ │ │
│ 2040 │ │ if len(error_msgs) > 0: │
│ ❱ 2041 │ │ │ raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format( │
│ 2042 │ │ │ │ │ │ │ self.__class__.__name__, "\n\t".join(error_msgs))) │
│ 2043 │ │ return _IncompatibleKeys(missing_keys, unexpected_keys) │
│ 2044 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: Error(s) in loading state_dict for T5ForConditionalGeneration:
size mismatch for encoder.block.0.layer.0.SelfAttention.q.weight: copying a param with shape
torch.Size([524288, 1]) from checkpoint, the shape in current model is torch.Size([1024, 1024]).
size mismatch for encoder.block.0.layer.0.SelfAttention.k.weight: copying a param with shape
torch.Size([524288, 1]) from checkpoint, the shape in current model is torch.Size([1024, 1024]).
`

Does anyone have any solutions to this?