DataLoader fails with custom collate_fn

vision
4

you can try something like this:

import torch
from torch.utils.data import Dataset, DataLoader


class VarSizeImages(Dataset):
    
    def __len__(self):
        return 100
    
    def __getitem__(self, idx):
        size = torch.randint(10, 20, size=(2, ), dtype=torch.long)
        img = torch.rand(*size).numpy()
        label = torch.randint(0, 10, size=(1, ), dtype=torch.long).item()
        return img, label
        

ds = VarSizeImages()

dp = ds[0]
dp[0].shape, dp[1]

loader = DataLoader(ds, batch_size=10, collate_fn=lambda batch: [(torch.from_numpy(dp[0]), torch.tensor(dp[1])) for dp in batch])

for batch in loader:
    pass

print(len(batch), type(batch))
print(len(batch[0]), type(batch[0][0]), type(batch[0][1]))
> 10 <class 'list'>
> 2 <class 'torch.Tensor'> <class 'torch.Tensor'>

HTH