@mrshenli
Yeah. I have tried what you said. It didn’t work. Current setting:
def train(args):
    # Env
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.local_rank)
    torch.backends.cudnn.benchmark = True
    torch.cuda.set_device(args.local_rank)
    dist.init_process_group('nccl')
    device = torch.device('cuda', args.local_rank)