Is there a way to do the opposite: a model trained on a single GPU being loaded into a multi-GPU inference script? This is how I’ve been doing it, but it doesn’t work as expected:
    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)
    print("Using", torch.cuda.device_count(), "GPUs")

    model.to(device)
    checkpoint = torch.load(args.pretrained_path)

    if torch.cuda.device_count() > 1:
        model.module.load_state_dict(checkpoint['model_state_dict'])
    else:
        model.load_state_dict(checkpoint['model_state_dict'])

I’m using 2 GPUs, I’ve set the batch size to 50 (so 25 samples per GPU). The input right before being passed into the model has dimensions torch.Size([50, 3, 2048, 2048]). However, the output returned by the model only has 25 items. Not sure why this is the case. Appreciate any insights into this!