How to use lightning with a dataset that is generated on the GPU?

[turian]

I am a long-time lightning fan, on the other hand in the past I have gotten bitten by how opinionated it is about training loops. I have a new project in pure torch, and I'm debating whether I can easily migrate it to lightning. I have a related question here: #14229

My datatset is generated ON the GPU using torchsynth. All I need to know is the batch number, and that is used to deterministically generate the data on GPU. What is the best way to use lightning in this setting, to avoid unnecessary CPU-to-GPU data moves?

[rohitgr7]

if it's already on GPU, you can override the transfer_batch_to_device hook.

def transfer_batch_to_device(self, batch, *args, **kwargs):
     return batch

but note that it might create issues with the mulit-GPU case if your data is not on the same device as the model.