How to put all but some vars to GPU
By default, in Lightning everything that is returned by a dataset is collated by the data loader and shipped to the same device.
However, I am frequently in the situation where I have let's say x, y which are tensors and something like y_semantic which is in principle related to y but of higher data type, say a dictionary with some meta information about augmentations or so.
I don't need 'y_semantic' to be on the GPU. Is there some flag or some method that I can overwrite so that some variables stay on CPU?
Something like
def ship_batch(self, batch):
    batch[0] = batch[0].to(self.device)
    # ...
    batch[2] = batch[2].cpu() # just for illustration here
Dear @Haydnspass,
You have several ways to do this:


Create a custom Data / Batch Object and implement the .to function to move only what is required.


Simpler: Override LightningModule.transfer_batch_to_device hook and add your own logic to move only x, y to the right device.


Best,
T.C