Why DataLoader return list that has a different length with batch_size
I am writting a customed dataloader, while the returned value makes me confused.
import torch
import torch.nn as nn
import numpy as np
import torch.utils.data as data_utils

class TestDataset:
    def __init__(self):
        self.db = np.random.randn(20, 3, 60, 60)

    def __getitem__(self, idx):
        img = self.db[idx]
        return img, img.shape[1:]

    def __len__(self):
        return self.db.shape[0]


if __name__ == '__main__':
    test_dataset = TestDataset()
    test_dataloader = data_utils.DataLoader(test_dataset,
                                       batch_size=1,
                                       num_workers=4,
                                       shuffle=False, \
                                       pin_memory=True
                                       )
    for i, (imgs, sizes) in enumerate(test_dataloader):
        print(imgs.size())  # torch.Size([1, 3, 60, 60])
        print(sizes)  # [tensor([60]), tensor([60])]
        break

Why &quot;sizes&quot; returns a list of length 2? I think it should be &quot;torch.Size([1, 2])&quot; which indicates height and width of a image(1 batch_size).
Further more, should the length of the returned list be the same to batch_size? If I want to get the size, I have to write &quot;sizes = [sizes[0][0].item(), sizes[1][0].item()]&quot;. And this makes me very confused.
Thank you for your time.

It is caused by collate_fn function and its default behaviour. It's main purpose is to ease the batch preparation process. So you can customize your batch preparation process updating this function. As stated in documentation collate_fn, it automatically converts NumPy arrays and Python numerical values into PyTorch Tensors and it preserves the data structure. So it returns in your case [tensor([60]), tensor([60])]. In many cases, you return image with labels as tensors(instead of size of image) and feedforward to neural net. I don't know why you return the image size while enumerating, but you can get what you need adding a custom collate_fn as:
def collate_fn(data):
    imgs, lengths = data[0][0],data[0][1]    
    return torch.tensor(imgs), torch.tensor([lengths])

Then you should set it to DataLoader's argument:
test_dataloader = DataLoader(test_dataset,
                                    batch_size=1,
                                    num_workers=4,
                                    shuffle=False, \
                                    pin_memory=True, collate_fn=collate_fn
                                    )

Then you can loop as:
for i, (imgs, sizes) in enumerate(test_dataloader):
    print(imgs.size())
    print(sizes)  
    print(sizes.size())  
    break

and output will be as:
torch.Size([3, 60, 60])
tensor([[60, 60]])
torch.Size([1, 2])

Afterall, I would like to add one more thing, you should not just return the self.db.shape[0] in len function. In this case your batch size is 1 and it's ok; however, when the batch size changes it will not return the true value for #batches. You can update your class as:
class TestDataset:
    def __init__(self, batch_size=1):
        self.db = np.random.randn(20, 3, 60, 60)
        self._batch_size = batch_size
        
    def __getitem__(self, idx):
        img = self.db[idx]
        return img, img.shape[1:]

    def __len__(self):
        return self.db.shape[0]/self._batch_size

