HeteroData and Dataloader

[jkalogero]

Hi,
I'm implementing a model, where each input will have 2 types of nodes (N1=36, N2=10) and 3 types of undirected edges.
I thought using HeteroData would be a good idea and created my Dataset as proposed in the documentation, but instead of saving a Data object each time, I saved an HeteroData object.
Then, I tried using a Dataloader, but on the classic loop:

for batch in train_dataloader:
    . . .

I get the error:
TypeError: DataLoader found invalid type: <class 'torch_geometric.data.hetero_data.HeteroData'>

After searching the documentation, I didn't find any other loader that would work with HeteroData, since I don't want any sampling on the graphs.

Is HeteroData not supported by the classic Dataloader or am I doing something wrong?
Any suggestion would be welcome.
Also thanks for the awesome library!

[rusty1s]

Ups, that should be supported. I fixed it in master :) Thanks for letting me know.