How to determine the device from within a LightningDataModule

[nathanpainchaud]

❓ Questions and Help

What is your question?

Is there a recommended way (or is it even at all possible) to determine which device is used from within the new LightningDataModule?

I ask the question because right now I decide how to set pin_memory when initializing the dataloaders based on the device. Prior to 0.9, the _dataloaders() methods were part of the LightningModule and I could simply access self.device to check the device the data would be transferred to.

Code

Example of what I did before 0.9:

def train_dataloader(self) -> DataLoader:
    return DataLoader(..., pin_memory=self.device.type == "cuda")

What have you tried?

I've checked the doc and API of the LightningDataModule, but I don't see any indications of how to get the device.

What's your environment?

• Version: 0.9.0

[ananyahjha93]

you can check the trainer object for on_gpu variable and override your datamodule method to take this as a param.

[nathanpainchaud]

For anyone wondering, the recent PR #3684, which adds a reference to the Trainer in the LightningDataModule, fixes my issue and allows us to do something like this:

def train_dataloader(self) -> DataLoader:
    return DataLoader(..., pin_memory=self.trainer.on_gpu)