Hi @ptrblck thanks for the prompt reply.
I am trying to distribute data across multiple GPUs, because I’ve found that training time decreases dramatically when I’ve allocated them on the GPU before I start training.
As you can imagine, with large data-sets that’s very hard to do,  the GPU memory is also used by the model and the batch samples. So I am looking into having parts of the data-set across multiple-GPUs.
According to the error produced, and as you explained that’s a problem because the weights, batch samples and other ops (CE Loss tensors) need to be on the same device?
It is my understanding (please correct me if I’m wrong) that I need to have copies of the network/model across multiple GPUs (which is where nn.DataParallel comes into play) and I also need to ensure that the batch tensors are also on the same GPU?
It is the last part I don’t know how is done.
I’ve been looking into nn.DistributedDataParallel but the tutorial doesn’t explicitly state how to distribute Dataloader across multiple GPU, and even more importantly, how to ensure the device id is the same as the model and batch during training. I am sure it is something very simple.
Alternatively, I was wondering if I should have multiple dataloaders each associated with a GPU, and then just train a model across those different dataloaders/GPUs.
#EDIT
To be more clear, I’ve been looking at the tutorial here: https://pytorch.org/tutorials/intermediate/ddp_tutorial.html 49
The nn.DistributedDataParallel example explains how to send the model and other parameters across seamlessly:
    model = some_model().to(device_ids[0])
    ddp_model = DDP(model, device_ids=device_ids)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(device_ids[0])
    loss_fn(outputs, labels).backward()
    optimizer.step()

What however I don’t know, is how to ensure that the batch input (in this case the torch.randn) is on the same cuda device. For example, the tutorial from nn.DataParallel (https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html 25) shows that within the epoch and batch training loop, I copy (or upload) to a device:
for data in rand_loader:
    input = data.to(device)
    output = model(input)

Suppose the above (^^^) is my training loop, and my tensors (input and associated output/classes) have been already loaded to different GPUs, how can I ensure that they are on the correct gpu device?
The closest I’ve seen to explaining how to do this is this tutorial (https://pytorch.org/tutorials/beginner/former_torchies/parallelism_tutorial.html 24) but it doesn’t actually answer that question.
Many thanks!