`logits_py` should be the log of the current `logits_py` or we can just give it as `probs` and not `logits` to the corresponding distribution.
logits_pyshould be the log of the currentlogits_pyor we can just give it asprobsand notlogitsto the corresponding distribution.