@ptrblck, turns out that doesn’t quite work.
Right now, my getitem fn does a lot of useless I/O and preprocessing. It gets the N K C H W tensor which the collate function you wrote will reduce to N C H W. (Shown on the left in my diagram)

image.jpg1251×938 180 KB

What I want is shown in the diagram on the right. getitem returns a 1 C H W tensor (no useless I/O and preprocessing). The 1 here is got by sampling an index from 0 to K, and it has to be the same across the batch. (I also need to know which index it is…)
Does that make sense? Any ideas on how to proceed?
Maybe this is too specific a question to my usecase…
Thanks for your help!