I want grad to exactly contain the stack of gradients wrt to three losses. How I am doing it at present is this way:
    flattened_grads = [torch.Tensor([]).to(device) for i in range(len(losses))]
    
 

    # Accumulate the gradient vectors
    for i, loss in enumerate(losses):
        for optim in optimizers:
            optim.zero_grad()
        loss.backward(retain_graph = True)
        for net in nets:
            for name, param in net.named_parameters():
                flattened_grads[i] = torch.cat([flattened_grads[i],param.grad.view(-1)])

I want to optimize this code to run faster and possibly in parallel. As the computation graphs are similar I should ideally backprop each node once.