diff --git a/hypernerf/model_utils.py b/hypernerf/model_utils.py
index 6528297..7e3cbea 100644
--- a/hypernerf/model_utils.py
+++ b/hypernerf/model_utils.py
@@ -111,7 +111,7 @@ def volumetric_rendering(rgb,
   last_sample_z = 1e10 if sample_at_infinity else 1e-19
   dists = jnp.concatenate([
       z_vals[..., 1:] - z_vals[..., :-1],
-      jnp.broadcast_to([last_sample_z], z_vals[..., :1].shape)
+      jnp.broadcast_to(jnp.array([last_sample_z]), z_vals[..., :1].shape)
   ], -1)
   dists = dists * jnp.linalg.norm(dirs[..., None, :], axis=-1)
   alpha = 1.0 - jnp.exp(-sigma * dists)