I want to make a network that has three parts - the first one is choosing the next one to be used:
model = nn.Sequential()
# hinge is only used to choose between subnetwork_1 and _2
model.add_module('0: hinge', nn.Linear(3, 2))
# based on the choice made by hinge, one of the following is used: 
model.add_module('1: subnetwork_1', nn.Linear(3, 1))
model.add_module('2: subnetwork_2', nn.Linear(3, 1))

x = torch.randn((1, 3), requires_grad=True)

z = model[0](x)
_, max_indice = torch.max(z, 1)

x = model[max_indice+1](x)

[Google Colab](Code to play is available on Google Colab)
The problem is that the layer ‘hinge’ does not learn anything, so my network does not become better and better in choosing the best subnetwork for further processing the data. I already know [max_indice] will not work, but what will?..