Fix recompliation of model

ecmwf-lab · Dec 10, 2024 · 104096d · 104096d
1 parent 18e7458
commit 104096d
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/ai_models_gencast/model.py b/src/ai_models_gencast/model.py
@@ -217,7 +217,9 @@ def run_forward(
                 LOG.info("Model license: %s", self.ckpt.license)
 
             jax.jit(self._with_configs(run_forward.init))
-            self.model = self._drop_state(self._with_params(jax.jit(self._with_configs(run_forward.apply))))
+            self.model = xarray_jax.pmap(
+                self._drop_state(self._with_params(jax.jit(self._with_configs(run_forward.apply)))), dim="sample"
+            )
 
     def run(self):
 
@@ -290,7 +292,7 @@ def run(self):
 
             chunks = []
             for chunk in rollout.chunked_prediction_generator_multiple_runs(
-                xarray_jax.pmap(self.model, dim="sample"),
+                self.model,
                 rngs=rngs,
                 inputs=input_xr,
                 targets_template=template * np.nan,