PyTorch: пользовательский поэлементный слой

Question

Я пытаюсь создать слой в PyTorch, который демонстрирует некоторое пользовательское поведение.

Мой текущий подход выглядит довольно стандартным

class MyElementwiseLayer(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(layer):
        for neuron in layer:
            neuron = my_function(neuron)
        return layer

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

Я предполагаю, например, что когда вы применяете nn.relu, torch может эффективно выполнять эти вычисления для всех доступных ресурсов. Поскольку у меня есть 14 ядер для игры, мне интересно, как можно использовать некоторые из врожденных способностей факелов, не выполняя какой-либо хакерский параллелизм сам.