add dtype to kernel (#365)

ott-jax · May 25, 2023 · 64d26ac · 64d26ac
1 parent 795fabe
commit 64d26ac
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/ott/solvers/nn/layers.py b/src/ott/solvers/nn/layers.py
@@ -65,6 +65,7 @@ def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
         "kernel", self.kernel_init, (inputs.shape[-1], self.dim_hidden)
     )
     kernel = self.rectifier_fn(kernel)
+    kernel = jnp.asarray(kernel, self.dtype)
     y = jax.lax.dot_general(
         inputs,
         kernel, (((inputs.ndim - 1,), (0,)), ((), ())),