The docs for gather describe how its indexing works: https://pytorch.org/docs/stable/torch.html#torch.gather 128.
In particular, it sounds like you want something like:
out[i][j][k] = input[i][indices[i]][k]

We can massage this a little into the form that torch.gather wants, which is out[i][j][k] = input[i][index[i][j][k]][k].
What happens if you do the following:
indices = indices.unsqueeze(-1).unsqueeze(-1)
# indices is now size (16, 1, 1). Let's turn it into the same size as action_prop
indices = indices.expand_as(action_prob)
action_prob.gather(1, indices)