Attribute Error when trying to run SVI

I am getting:
AttributeError: DynamicJaxprTracer has no attribute loss_with_mutable_state
when trying to run SVI. I am running the model in a Kaggle notebook copied from another Kaggle notebook where the model worked fine. Googling doesn’t give anything. Any ideas for solving this?

Without further information, it is hard to tell. You can check if your svi construction matches the tutorials.