Transferring Data Augmentation Modules to Device

[EthanMarx]

Hello! In my DataModule I setup a couple of torch.nn.Modules that do on-gpu data augmentations. I am wondering how to ensure that these Modules get transferred to the appropriate device when doing multi gpu training?

My setup looks something like the pseudo code below

class DataModule(pl.DataModule):
    
    def __init__(self):
        # define some torch.nn.Module
        #  for on-gpu augmentations
        self.augmentor = Augmentor()

    def train_dataloader(self):
        # dataloader loads data from disk
        dataloader = ...
        return dataloader

    def on_after_batch_transfer(self, batch, _):
        # augment batch once on gpu. 
        # how do I ensure augmentor is on correct device?
        batch = self.augmentor(batch)

appreciate the help