Send the output of your network to DataLoader

My training regime is set up so that the next batch depends on the output of the network on the previous batch. I’m currently saving the output on the previous batch on the model itself to be used for the next batch, but it doesn’t seem like the correct way to go about this, and I was wondering if there’s a way to send the output back to DataLoader so that I can keep the forward pass simpler. My initial thought was to use IterableDataset and send the output back via Python generator, but this doesn’t seem possible based on my reading of the DataLoader code as there is no way to access the generator created by IterableDataset. Is there a way to do this? Is my current solution of saving the output on the previous batch the best one, or is there a better solution?