diff --git a/recml/inference/models/jax/DLRM_DCNv2/dlrm_model.py b/recml/inference/models/jax/DLRM_DCNv2/dlrm_model.py
index 23720e5..a38c4c4 100644
--- a/recml/inference/models/jax/DLRM_DCNv2/dlrm_model.py
+++ b/recml/inference/models/jax/DLRM_DCNv2/dlrm_model.py
@@ -135,7 +135,7 @@ def __call__(
 
     sparse_embeddings = embed.SparseCoreEmbed(
         feature_specs=self.feature_specs,
-        mesh=self.mesh,
+        _mesh=self.mesh,
         sharding_axis=self.sharding_axis,
     )(embedding_lookups)
     sparse_embeddings = jax.tree.flatten(sparse_embeddings)