Using unmodified input

Hello, I am trying to learn two distributions with my network. The first one is a categorical distribution and the second one a normal distribution. Below is my pseudocode of the forward() function which illustrates the problem. From input2 i only take the value at that index, which the categorical distribution outputs and then use this value as the mean of the normal distribution. But autograd returns Nones since somehow the output is not connected to input. I hope there is a workaround, for which I would be very thankful! The action_std is a parameter defined in init.

def forward(input1, input2):

x=softmax(layer1(input1))
dist1 = Canonical(x)
index = dist1.sample()
action_mean = input2[index]
dist2 = Normal(action_mean, action_std)
action = dist2.sample()

return action