Returning a Class from a Databunch Instead of a Tensor

edit: resolved. The class I was dealing with had a special version of to_device that didn’t return anything. This was causing the issue

I’ve been playing around with Deep Graph Library and I’m trying to create a dataloader for graph objects.

In DGL, your basic data object is a DGLGraph class object that contains node and edge data and has methods to handle message passing, gathering and other functions. This is the actual object that must be fed to the graph network, along with a feature matrix for the graph nodes.

I’m able to create custom ItemBase/ItemList classes for the graphs and get them into a Databunch. I’m also able to create standard Pytorch dataloaders and create a Databunch from those. Batching is handled with a custom collate function.

The collate function returns data as two tuples. The input data is a tuple of (batch_graph, feature_vector) where batch_graph is a DGLGraph object and feature_vector is a pytorch tensor. The target data is a tuple of (label_vector, mask), both of which are pytorch tensors.

When I pull a batch from the Databunch, the feature vector, label vector and mask are returned fine. However the DGLGraph objects are returned as None. I think this is caused by returning an object instead of a tensor, but I’m not sure. I’ve verified that the object is created correctly in the collate function - the issue must be somewhere after that.

Does anyone know what function in Databunch I would need to change to fix this?