diff --git a/init2winit/model_lib/dlrm.py b/init2winit/model_lib/dlrm.py
index 888d9f3d..49d4a68b 100644
--- a/init2winit/model_lib/dlrm.py
+++ b/init2winit/model_lib/dlrm.py
@@ -22,6 +22,7 @@
 from init2winit.model_lib import model_utils
 import jax
 from jax import nn as jnn
+import jax.experimental.shard_map
 import jax.numpy as jnp
 from ml_collections.config_dict import config_dict