run dgllife pubchem_aromaticity example report error

Question

I downloaded dgllife and run the pubchem_aromaticity example. (https://github.com/chaoyue729/dgl-lifesci/tree/master/examples/property_prediction/pubchem_aromaticity). But it's always report error. When I change the args['device']="cpu" it can run. But it's too slow. I need run it on cuda. How can I fix it?

def main(args):
    args['device'] = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
    #args['device'] = torch.device("cpu")
    ......

---

dgl._ffi.base.DGLError: Cannot assign node feature "hv" on device cuda:0 to a graph on device cpu. Call DGLGraph.to() to copy the graph to the same device.

---

I guess the reason for the error is main.py's bg in line 46, whose type is "dgl.heterograph.DGLHeteroGraph", cannot be copied to CUDA. Reference (https://docs.dgl.ai/guide_cn/graph-gpu.html?highlight=dglerror). But I don't know how to set it.

Answer 1

I have solved this problem. The solution is add some code on the function regress by "main.py".

def regress(args, model, bg):
    atom_feats, bond_feats = bg.ndata.pop('hv'), bg.edata.pop('he')
    atom_feats, bond_feats = atom_feats.to(args['device']), bond_feats.to(args['device'])
    bg = bg.to(args['device'])
    return model(bg, atom_feats, bond_feats)