Использование torch.nn.DataParallel с пользовательским расширением CUDA

Насколько я понимаю, встроенные операции pytorch автоматически обрабатывают партии посредством неявной векторизации, позволяя параллелизм между несколькими GPU.

Однако при написании пользовательской операции в CUDA в соответствии с Документацией приведенный пример LLTM выполняет операции, которые являются пакетными, например, вычисляя градиент функции Sigmoid по-разному.

Тем не менее, у меня есть вариант использования, который не является периодическим и не является векторизованным. Находясь на одном графическом процессоре, я в настоящее время (неэффективно) перебираю каждый элемент в пакете, выполняя запуск ядра для каждого, например (написанный в браузере, чтобы продемонстрировать):

std::vector<at::Tensor> op_cuda_forward(at::Tensor input, 
                                        at::Tensor elementSpecificParam) {

    auto output = at::zeros(torch::CUDA(/* TYPE */), {/* DIMENSIONS */});

    const size_t blockDim = //
    const size_t gridDim = //
    const size_t = numBatches = //

    for (size_t i = 0; i < numBatches; i++) {
        op_cuda_forward_kernel<T><<<gridDim, blockDim>>>(input[i],
                                                         elementSpecificParam[i], 
                                                         output[i]);
    }

    return {output};
}

Однако я хочу разделить эту операцию на несколько графических процессоров с помощью пакетного элемента.

Как распределение output тензора будет работать в сценарии с несколькими GPU?

Конечно, можно создать промежуточные тензоры на каждом графическом процессоре, прежде чем запускать соответствующее ядро, однако накладные расходы на копирование входных данных на каждый графический процессор и обратно будут проблематичными.

Есть ли более простой способ запуска ядра без предварительного тестирования среды для информации о GPU (# GPU и т.д.)?

Конечной целью является работа CUDA, которая работает с torch.nn.DataParallel.