How to move data to the cuda in customized datacollator in DDP mode

[yuzhenmao]

Hi, for some reason I need to move a batch to gpu in my own datacollator function within the dataloader. Is it possible to do this in the ddp setting by pytorch-lightning? Or do I need to manually move the data by myself using “local_rank = os.environ.get(‘LOCAL_RANK’) & data.to(device)”? Thanks!