Dataloader returns a list of batch_size

So in my dataset, each sample has the different first dimension, for example x1.shape = (2, 100) and y1.shape = (2, 1), and sample 2 could be x2.shape = (5, 100) and y2.shape = (5,1)

How could I customize my collate_fn so that I could return a loader list of batch_size. For example, if my batch size = 2, then how could I get:
for inputs , labels in train_loader:
# inouts is a list of length 2, with first element ( 2,100), second element (5, 100)
# labels is a list of length 2, with first element ( 2,1), second element (5, 1)

If I understood your question correctly, then the answer would be: by simply returning x[n] and y[n] from your Dataset class. Something like the following:

import torch

class Dataset(torch.nn.Module):
    def __init__(self):
    def __len__(self):
        return 10
    def __getitem__(self, ind):
        inputs = torch.rand(2)
        labels = torch.rand(2)
        return inputs, labels

dataset = Dataset()
data_loader =, batch_size=2)
for inputs, labels in data_loader:
    print("Inputs: ", inputs)
    print("Labels: ", labels)