I have a similar issue, which I managed to solved by changing the arrangement of the data, but I cannot understand why it could not work before. So I would like someone to explain it to me if possible. This is my initial dataset class:
class ModelDataset(Dataset):
    
    def __init__(self, data, labels, ind):

        # take only the data specified by ind. ind is a list of indexes
        self.data = data[:,ind,:]  # size = (6, ind, 256)
        self.labels = labels[ind]
       
    def __getitem__(self, index):
        
        pair = self.data[:,index,:] # pair = (6,256)
        label = self.labels[index]
        
        x1 = pair[:, :128]
        x2 = pair[:, 128:]
        
        return x1.T, x2.T, label
    
    def __len__(self):
        return len(self.data)
    

So when I was creating a dataloader on this dataset, even though I was specifing the batch size at a certain value the maximum batch size was stuck at 6. I thought that maybe the 3D arrangement of the self.data was the issue so I rearranged the self.data to be a 2D matrix and after that I was able to get the batch size I wanted. But I cannot understand why I could not make it work before. When I was fixing the batch size to be lower than 6 I was getting the correct data from the dataloader. I checked that.