Compatibility when DataLoader returns multiple batches, for prefetching purposes
I'm trying to convert https://github.com/xcmyz/FastSpeech to Pytorch Lightning. The code does something complicated in the DataLoader where batch_size**2 batches are placed on the device, then iterated over, effectively prefetching batch_size batches to use in an inner loop of your typical pytorch training loop (here)
Is there an easy way to make this compatible with Pytorch Lightning without giving up this prefetching optimization?
Edit: I think I was incorrect here, it's not prefetching, he's just defined a DataLoader that returns batches of batches. I think I can probably fix this.