Data processing as a batch way

34

Nevermind, found a solution here! Dataset Multiple Samples per getitem Call - #2 by ptrblck

Using a custom collate_fn would be more elegant though!