Error while fine-tuning model using example from "examples/finetune.ipynb": `RuntimeError: shape '[8, 129, 12, 512]' is invalid for input of size 8306688`

[BradyBonnette]

Hi! I've been experimenting with this model for a few things, and so far I like where it's going.

I want to attempt some fine-tuning, so I followed the same notebook found here: https://github.com/urchade/GLiNER/blob/main/examples/finetune.ipynb

However, when it comes to the training step, I am seeing the following error:

RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/parallel/parallel_apply.py", line 83, in _worker
    output = module(*input, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/gliner/model.py", line 103, in forward
    output = self.model(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/gliner/modeling/base.py", line 232, in forward
    span_rep = self.span_rep_layer(words_embedding, span_idx)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/gliner/modeling/span_rep.py", line 356, in forward
    return self.span_rep_layer(x, *args)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/gliner/modeling/span_rep.py", line 286, in forward
    return self.out_project(cat).view(B, L, self.max_width, D)
RuntimeError: shape '[8, 129, 12, 512]' is invalid for input of size 8306688

I tried to debug it myself, but I am rather confused as to why this is happening. All previous steps are copied verbatim from the example finetune notebook.

Maybe some useful info:

torch 2.4.0+cu121
gliner 0.2.10
transformers 4.45.0.dev0
huggingface_hub 0.23.2
accelerate 0.34.0
Python 3.10.12

Any ideas? I can provide more information where needed.

Thanks!