Dataloader custom collate for different input sizes

Have a look at this approach posted by @vfdev-5 and let us know, if that helps or not.