add back decorator for computing pareto fronts

Lisa · Lisa · commit 9911e744aa89 · 2025-04-04T16:12:20.000+01:00
diff --git a/qdax/core/containers/mome_repertoire.py b/qdax/core/containers/mome_repertoire.py
@@ -418,6 +418,7 @@ def init(  # type: ignore
 
         return new_repertoire  # type: ignore
 
+    @jax.jit
     def compute_global_pareto_front(
         self,
     ) -> Tuple[ParetoFront[Fitness], Mask]:
