How to tell Lightning which data to load onto the GPU (for 3rd party compatibility)
Hi !
I'm currently working on a segmentation network using PyTorch Lightning and MONAI.
Context

In my LightningDataModule, I preprocess my images by applying transforms (e.g., to resample them) before feeding them to my DataLoaders. For the transforms, I use a 3rd party framework called MONAI. It stores the context information of all the applied transforms in a structured nested dictionary.
After predicting the labels of my predict dataset, I would like to invert these transforms (e.g., to recover the original pixel dimensions) in my predict_step.

Problem

MONAI's inverting logic requires the context information to be numpy data or CPU tensors (GPU tensors are not supported). However, by default, Lightning moves all the data of my batch (including its context info) to GPU. It causes the bug I reproduced in this test example notebook (cf. last cell).

My question
How can I tell Lightning to only move the images and their labels to GPU and keep the context info on the CPU?
For further details, please refer to the following discussion Project-MONAI/MONAI#2348.
Thanks for your help !
Hi,
your LightningModule has a hook def transfer_batch_to_device(self, batch: Any, device: torch.device, dataloader_idx: int) -> Any: that you can override and which should be a perfect fit for that. Just make sure to only use the provided device.