Neural network 将torch.nn.DataParallel与自定义CUDA扩展一起使用

据我所知，内置的PyTorch操作都通过隐式矢量化自动处理批处理，允许跨多个GPU并行

然而，当根据在CUDA中编写自定义操作时，给出的LLTM示例执行批处理不变的操作，例如计算Sigmoid函数的元素梯度

然而，我有一个用例，它不是批处理元素不变的，也不是矢量化的。在单个GPU上运行，我目前（低效地）循环批处理中的每个元素，为每个元素执行内核启动，就像这样（写在浏览器中，只是为了演示）：

std:：vector

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};
}