Cool!
One problem that could come up here is that the sampler might get stuck in one of the local minima. So I would add a traceplot and also run several chains and check the rhat (pm.gelman_rubin). Our current hamiltonian methods don’t deal with multimodal posteriors particularly well. SMC might help with that, but I don’t know much about it.
Also, with nuts you don’t usually need that many samples (that is, as long as the sampler doesn’t get stuck around one mode for long time)