def my_collate(batch):
data = [item[0] for item in batch if item[0].shape[1]==256 and item[0].shape[2]==256]
target = [item[1] for item in batch if item[0].shape[1]==256 and item[0].shape[2]==256]
return [data, target]
I tried creating a custom collate_fn for my dataloader such that the stack contains only images of [3,256,256] shape.Can anyone please tell me why i am facing KeyError 0 ?
It’s Urgent.
Check what batch and item contains as it seems you cannot index the latter with [0].