Hi discordiaZA, I haven’t tried writing a custom distribution yet and decided to start by making the choice deterministic, using p_high.argmax() like junpenglao suggested. It is not working yet (see here), but I’m hoping that I’ll get it to work soon. What is the issue with your model?