Using scale in numpyro

That worked. Thanks!