How to put all but some vars to GPU

[Haydnspass]

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

[tchaton]

Dear @Haydnspass,

You have several ways to do this:

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

2. 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