Stratify the Data list of Pytorch_geometric.data.Data without creating Dataset

[naveen-kinnal]

Hello. I have a custom dataset for which I am building a list of Data objects

def build_gnn_dataset(sentences, labels):
    data_list = []
    for i in range(len(sentences)):
        ### MY CODE ###

        data = Data(x=x, y=y, edge_index=torch.t(edge_index))
        data_list.append(data)
    return data_list

I am directly feeding this to the DataLoader (without creating a Dataset object), since I need not save my data

data_loader = torch_geometric.data.DataLoader(data_list, batch_size=64, shuffle=True)

My custom dataset contains labels in 0,1 and is the class distribution is not proper. Hence, I want to stratify my batches in order to equally distribute the labels.

How can I do this directly using the Data List?

[wsad1]

In general you could pass a sampler to your DataLoader. Check out this example.