DGLError - Githubissues

I get this error when running the predict.py code, this is on an Ubuntu Linux 22.04 with an NVIDIA RTX A6000 GPU.
What could be the issue? Could it be that this is using the wrong version of some python or CUDA libraries?
Thanks in advance,
$ lsb_release -a
No LSB modules are available.
Distributor ID: Ubuntu
Description:    Ubuntu 22.04.3 LTS
Release:        22.04
Codename:       jammy

$ dpkg -l nvidia-driver-535
Desired=Unknown/Install/Remove/Purge/Hold
| Status=Not/Inst/Conf-files/Unpacked/halF-conf/Half-inst/trig-aWait/Trig-pend
|/ Err?=(none)/Reinst-required (Status,Err: uppercase=bad)
||/ Name              Version                     Architecture Description
+++-=================-===========================-============-=================================
ii  nvidia-driver-535 535.129.03-0ubuntu0.22.04.1 amd64        NVIDIA driver metapackage

$ dpkg -l nvidia-cuda-toolkit
Desired=Unknown/Install/Remove/Purge/Hold
| Status=Not/Inst/Conf-files/Unpacked/halF-conf/Half-inst/trig-aWait/Trig-pend
|/ Err?=(none)/Reinst-required (Status,Err: uppercase=bad)
||/ Name                Version         Architecture Description
+++-===================-===============-============-=================================
ii  nvidia-cuda-toolkit 11.5.1-1ubuntu1 amd64        NVIDIA CUDA development toolkit

Error:

N=1 L=491
Traceback (most recent call last):
  File "/data2/cynicjon/RoseTTAFold2/network/predict.py", line 493, in <module>
    pred.predict(
  File "/data2/cynicjon/RoseTTAFold2/network/predict.py", line 318, in predict
    self.run_prediction(
  File "/data2/cynicjon/RoseTTAFold2/network/predict.py", line 392, in run_prediction
    logit_s, logit_aa_s, _, logits_pae, p_bind, xyz_prev, alpha, symmsub, pred_lddt, msa_prev, pair_prev, state_prev = self.model(
  File "/home/petmedix/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data2/cynicjon/RoseTTAFold2/network/RoseTTAFoldModel.py", line 86, in forward
    msa, pair, R, T, alpha, state, symmsub = self.simulator(
  File "/home/petmedix/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data2/cynicjon/RoseTTAFold2/network/Track_module.py", line 659, in forward
    msa_full, pair, R_in, T_in, state, alpha, symmsub = self.extra_block[i_m](msa_full, pair,
  File "/home/petmedix/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data2/cynicjon/RoseTTAFold2/network/Track_module.py", line 562, in forward
    R, T, state, alpha = self.str2str(
  File "/home/petmedix/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/petmedix/.local/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 14, in decorate_autocast
    return func(*args, **kwargs)
  File "/data2/cynicjon/RoseTTAFold2/network/Track_module.py", line 500, in forward
    shift = self.se3(G, node, l1_feats, edge_feats)
  File "/home/petmedix/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data2/cynicjon/RoseTTAFold2/network/SE3_network.py", line 89, in forward
    return self.se3(G, node_features, edge_features)
  File "/home/petmedix/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/se3_transformer-1.0.0-py3.10.egg/se3_transformer/model/transformer.py", line 185, in forward
  File "/home/petmedix/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/se3_transformer-1.0.0-py3.10.egg/se3_transformer/model/transformer.py", line 47, in forward
  File "/home/petmedix/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/se3_transformer-1.0.0-py3.10.egg/se3_transformer/model/layers/attention.py", line 166, in forward
  File "/home/petmedix/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/se3_transformer-1.0.0-py3.10.egg/se3_transformer/model/layers/convolution.py", line 319, in forward
  File "/usr/local/lib/python3.10/dist-packages/dgl/view.py", line 179, in __call__
    return self._graph.all_edges(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/dgl/heterograph.py", line 3589, in all_edges
    src, dst, eid = self._graph.edges(self.get_etype_id(etype), order)
  File "/usr/local/lib/python3.10/dist-packages/dgl/heterograph_index.py", line 696, in edges
    edge_array = _CAPI_DGLHeteroEdges(self, int(etype), order)
  File "dgl/_ffi/_cython/./function.pxi", line 295, in dgl._ffi._cy3.core.FunctionBase.__call__
  File "dgl/_ffi/_cython/./function.pxi", line 227, in dgl._ffi._cy3.core.FuncCall
  File "dgl/_ffi/_cython/./function.pxi", line 217, in dgl._ffi._cy3.core.FuncCall3
dgl._ffi.base.DGLError: [11:50:51] /opt/dgl/src/array/array.cc:42: Operator Range does not support cuda device.
Stack trace:
  [bt] (0) /usr/local/lib/python3.10/dist-packages/dgl/libdgl.so(+0xf9f6f) [0x7f53c24f9f6f]
  [bt] (1) /usr/local/lib/python3.10/dist-packages/dgl/libdgl.so(dgl::aten::Range(long, long, unsigned char, DGLContext)+0x117) [0x7f53c24fb1b7]
  [bt] (2) /usr/local/lib/python3.10/dist-packages/dgl/libdgl.so(dgl::UnitGraph::COO::Edges(unsigned long, std::string const&) const+0x9b) [0x7f53c2a17bbb]
  [bt] (3) /usr/local/lib/python3.10/dist-packages/dgl/libdgl.so(dgl::UnitGraph::Edges(unsigned long, std::string const&) const+0x9e) [0x7f53c2a1095e]
  [bt] (4) /usr/local/lib/python3.10/dist-packages/dgl/libdgl.so(dgl::HeteroGraph::Edges(unsigned long, std::string const&) const+0x2a) [0x7f53c29112ca]
  [bt] (5) /usr/local/lib/python3.10/dist-packages/dgl/libdgl.so(+0x51ab51) [0x7f53c291ab51]
  [bt] (6) /usr/local/lib/python3.10/dist-packages/dgl/libdgl.so(DGLFuncCall+0x48) [0x7f53c28a9598]
  [bt] (7) /usr/local/lib/python3.10/dist-packages/dgl/_ffi/_cy3/core.cpython-310-x86_64-linux-gnu.so(+0x15413) [0x7f53c1015413]
  [bt] (8) /usr/local/lib/python3.10/dist-packages/dgl/_ffi/_cy3/core.cpython-310-x86_64-linux-gnu.so(+0x15c2b) [0x7f53c1015c2b]
uw-ipd / RoseTTAFold2

DGLError #24
