expose device in dataset / DataModule

[slbuilder]

Is there a way to access the device being used for the instance of DataModule/Dataset.

My requirement is to be able to load data directly on GPU without the copy overhead of cpu to gpu. I already have the code to do it. I specify and pass the device around in native PyTorch like so.

class MyDataset(torch.utils.data.Dataset):
      def __init__(self, args):
           self.device = torch.device(args.device)
      def __getitem__(self, idx):
          # load corresponding data on self.device
          # this does not create the overhead of loading the tensor on cpu first and then calling a to(self.device)
         return data

Is there a way to achieve something similar with the current device id exposed in Pytorch-lightning?