A similar problem on stackoverflow is here 4, but no answer is useful.
As this tutorial 4 shown, the output of multi-gpus will be concatenated on the dimension 0, but I don’t know why does it not work in my code.

    model = T2T(......)  # T2T is a sub class of nn.Module
    if torch.cuda.device_count() > 1:
        print("Using", torch.cuda.device_count(), "GPUs!")
        model = torch.nn.DataParallel(model)
    model = model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    train_loader = DataLoader(......)
    ......
    epoch_steps = len(train_loader)
    STEPS = train_epochs * epoch_steps
    train_loader = cycle(train_loader)
    ......
    for step in range(STEPS):
        batch_data = next(train_loader)
        optimizer.zero_grad()
        labels = batch_data['labels'].to(device)
        ......
        probs = model(......)
        # calculate loss using labels and probs

When calculating loss, I got an error, the batch size of labels and probs are not the same. The shape of labels is [64,…] and the shape of probs is [32,…]. I am using 2 gpus, so I guess the output of multi-gpus are not concatenated.
Any idea how to fix this? Thanks in advance!