nn.DataParallel (which seems to be used in your use case) could create an imbalanced memory usage and could thus cause an OOM on the default device, which is why we recommend to use DistributedDataParallel with a single process per GPU instead.
I assume you are also running OOM, so I would recommend to try to apply DDP.