Fix missing infer key in handlers.lift

[dominikstrb]

This PR addresses issue #891 and fixes handlers.lift in 0.5.0

[fehiepsi]

Thanks, @dominikstrb! Welcome to your first contribution! ;) Could you add a regression test for this at test/test_mcmc.py?

def test_model_with_lift_handler():
    def model(data):
        c = numpyro.param("c", jnp.array(1.), constraint=dist.constraints.positive)
        x = numpyro.sample("x", dist.LogNormal(c, 1.), obs=data)
        return x

    nuts_kernel = NUTS(numpyro.handlers.lift(model, prior={"c": dist.Gamma(0.01, 0.01)}))
    mcmc = MCMC(nuts_kernel, num_warmup=10, num_samples=10)
    mcmc.run(random.PRNGKey(1), x)

[dominikstrb]

I added the test. Thanks for your guidance!

[fehiepsi]

Awesome, thanks for the fix @Dominikstr!