What you can do is to designate a different GPU to do this. However, if the accumulated gradients fit on one GPU, it’s not worth thinking about hacks to do that on separate GPUs imho, because doing it on one device should be the fasted way to do it. Actually, I think they are accumulated to calculate the loss, and the gradients, I’ve sketched it out once …

dataparallel.png1275×2312 660 KB

The reason why it is computed on the default GPU is that you typically have something like this:
start_time = time.time()
for epoch in range(num_epochs):
    
    model.train()
    for batch_idx, (features, targets) in enumerate(train_loader):
        
        features = features.to(device)
        targets = targets.to(device)
            
        ### FORWARD AND BACK PROP
        logits, probas = model(features)
        cost = F.cross_entropy(logits, targets)
        optimizer.zero_grad()
        
        cost.backward()
        
        ### UPDATE MODEL PARAMETERS
        optimizer.step()

Note that targets (e.g., the class labels) is on the default GPU. Hence, if you want to do what you suggest you’d need to rewrite the code to split the class labels over the GPUs as well and then compute the loss on these separately and then combine the loss to compute the gradients. It’s not worth the hassle, imho, if the targets fit on the default GPU. Splitting it as well will only make it slower as the loss computation is super fast anyways.