Issue when using a cuda enabled installation

[Tonylac77]

Hi I am running into this issue with several torch/cuda versions :

Traceback (most recent call last):
  File "/home/tony/dockm8_packages/recipes/genscore_package/genscore/example/genscore.py", line 313, in <module>
    main()
  File "/home/tony/dockm8_packages/recipes/genscore_package/genscore/example/genscore.py", line 296, in main
    ids, scores = scoring(prot=inargs.prot,
  File "/home/tony/dockm8_packages/recipes/genscore_package/genscore/example/genscore.py", line 228, in scoring
    preds = run_an_eval_epoch(model,
  File "/home/tony/dockm8_packages/recipes/genscore_package/genscore/example/../GenScore/model/utils.py", line 269, in run_an_eval_epoch
    pi, sigma, mu, dist, atom_types, bond_types, batch = model(bgl, bgp)
  File "/home/tony/miniconda3/envs/genscore_2.3_gpu/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/tony/miniconda3/envs/genscore_2.3_gpu/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/tony/dockm8_packages/recipes/genscore_package/genscore/example/../GenScore/model/model.py", line 533, in forward
    C_batch = C_batch.repeat(1, N_l, N_t)[C_mask]
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)