Hi Guys,
For the case where you want to work with batches, and the image/label dimension in different batches might be different, you want to modify the collate_fn as pointed out by smth. The correct way is to construct a new collate fn.
Lets say that label for image 1 has dimension 10, whereas label for image 2 has dimension 15, and you are trying to create a batch of these images.
The main problem with unequal sizes is that the current collate_fn implementation uses itertools.zip. This creates an iteratable with the size of the smaller list, thus it will only create an iteratable with 10 elements in this case.
The solution is to use itertools.zip_longest. This will create an iteratable of dimension of the largest in the batch, i.e 15 in our cases. This function uses padding, so the values will be None for label 1 in the rest of the dimensions, which you will have to take care during run time in your code after you have got the batch.
Hope this helps! Let me know if anything else is needed