Does there exist a fast and convenient way for handling such a problem:
class MyModule(nn.Module):
def __init__(self):
super(MyModule, self).__init__()
self.linears = nn.ModuleList([nn.Linear(10, 20) for _ in range(10)])
def forward(self, x, indices):
x = self.linears[indices](x)
return x
You see i want to access different layers in the network conditioned on an additional input, which is also a list. Further i want to process the whole batch at once and the output.shape != input.shape.