Working code fails with PyMC 5.16

2

The error says that in the following line

choice_probs = pt.set_subtensor(
    choice_probs[all_subjs],
    pt.exp(beta[all_subjs] * belief[all_subjs,c])
     / (pt.exp(beta[all_subjs] * belief[all_subjs,0]) + pt.exp(beta[all_subjs] * belief[all_subjs,1]))
)

choice_probs[all_subjs] is a vector
And the update is a tensor3. Which one is correct?