我正在尝试在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火炬时，它可以在所有可用资源上有效执行该计算。由于我有14个核心可以玩，所以我想知道如何利用某些火炬的先天能力而不必自己实现一些hacky并发。