Dimension mismatch in `NUTS` and `metropolis`

Thanks for the update @mcavallaro!