CUDA support?

[davidress-ILW]

When using the EnriCo code base from GoogleLab, an error is being generated with the model running on CUDA when calling model.predict(). A modified stack trace is below:

Python: 3.10.14
torch version: 2.5.1+cu121

Code:

model = EnriCo.from_pretrained("urchade/enrico-beta").to('cuda')
model = model.eval()
out = model.predict(x, threshold=0.5, output_confidence=True)

Stacktrace:

File "C:\GraphERDemo\graphER_demo.py", line 127, in <module>
    r_list = process_text_and_relations(sent, [rel])
  File "C:\GraphERDemo\graphER_demo.py", line 60, in process_text_and_relations
    out = model.predict(x, threshold=0.5, output_confidence=True)
  File "C:\GraphERDemo\EnriCo\model.py", line 394, in predict
    out = self.forward(x, prediction_mode=True)
  File "C:\GraphERDemo\EnriCo\model.py", line 236, in forward
    filter_score_span, filter_loss_span = self._span_filtering(span_rep, x['span_label'])
  File "C:\Users\XXX\Analytics\lib\site-packages\torch\nn\modules\module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Users\XXX\Analytics\lib\site-packages\torch\nn\modules\module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\GraphERDemo\./EnriCo\modules\layers.py", line 111, in forward
    filter_score = filter_score.masked_fill(label == -1, -1e9)
RuntimeError: expected self and mask to be on the same device, but got mask on cpu and self on cuda:0

When using the GraphER code base available from the GitHub website, this also fails to use the CUDA device, giving a slightly different error.

Traceback (most recent call last):
  File "C:\GraphERDemo\graphER_demo.py", line 110, in <module>
    model = GraphER.from_pretrained("urchade/enrico-beta").to('cuda')
  File "C:\Users\XXX\_Analytics\lib\site-packages\huggingface_hub\utils\_validators.py", line 114, in _inner_fn
    return fn(*args, **kwargs)
  File "C:\Users\XXX\Analytics\lib\site-packages\huggingface_hub\hub_mixin.py", line 569, in from_pretrained
    instance = cls._from_pretrained(
  File "C:\GraphERDemo\./GraphER\modules\base.py", line 175, in _from_pretrained
    model = cls(config)
  File "C:\GraphERDemo\GraphER\model.py", line 70, in __init__
    num_heads=config.num_heads,
AttributeError: 'Namespace' object has no attribute 'num_heads'

Are there plans to support GraphER with CUDA? Or have I done something wrong? Do you have a GraphER with CUDA example?