wrt p0 is a non-differentiable operation, but you can use reinforcement learning, by using the reward(…) function on stochastic variables. so you will first get float indices, and you will stochastically sample from them over a multinomial or something, and you will give this sampled node a reward to propagate gradients back. see pytorch/examples or read on REINFORCE algorithm