Is it possible to permanently move a complete dataset to GPU memory before training?

[NiklasKappel]

I have a very small dataset on disk. I do not need to do any data augmentation. I would like to load the dataset into CPU memory, then move it to GPU memory, then proceed with training as usual. Is this possible and supported by lightning? What would be the best way to go about it?

In my mind, if I have only one GPU, the whole dataset should be in that GPU's memory. If I have multiple GPUs, one copy of the dataset should be in each GPU's memory.

In the single GPU case, I think I can move the dataset to GPU by calling .to("cuda") once on each tensor in the dataset. Then I would like to pass the dataset into a regular PyTorch dataloader and pass that dataloader to the lightning trainer's .fit method. Would that be correct?