fix default for intermediate quadratic layers in ICNN (#587)

ott-jax · Oct 9, 2024 · 0219671 · 0219671
1 parent 7309c76
commit 0219671
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/ott/neural/networks/icnn.py b/src/ott/neural/networks/icnn.py
@@ -116,7 +116,8 @@ def _get_wx(self, dim: int, rank: int) -> nn.Module:
         num_potentials=dim,
         use_linear=True,
         use_bias=True,
-        kernel_diag_init=nn.initializers.zeros,
+        kernel_diag_init=nn.initializers.constant(-2.0),
+        rectifier_fn=jax.nn.softplus,
         kernel_lr_init=self.init_fn,
         kernel_linear_init=self.init_fn,
         bias_init=nn.initializers.zeros,