last fixes.

ott-jax · Nov 14, 2022 · 09bffc0 · 09bffc0
1 parent 9ba0879
commit 09bffc0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ott/core/potentials.py b/ott/core/potentials.py
@@ -125,13 +125,13 @@ def _grad_g(self) -> Callable[[jnp.ndarray], jnp.ndarray]:
   @property
   def _grad_h_inv(self) -> Callable[[jnp.ndarray], jnp.ndarray]:
     assert isinstance(self.cost_fn, costs.TICost), (
-        "Cost must be RBF and ",
-        "provide access to Legendre Legendre transform of `h`."
+        "Cost must be a `TICost` and "
+        "provide access to Legendre transform of `h`."
     )
     return jax.vmap(jax.grad(self.cost_fn.h_legendre))
 
   def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]:
-    return [self._f, self._g, self.cost_fn], {"cor": self._cor}
+    return [self._f, self._g, self.cost_fn], {"corr": self._corr}
 
   @classmethod
   def tree_unflatten(