revert change

google · Oct 27, 2021 · b140571 · b140571
1 parent 240d974
commit b140571
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flax/linen/activation.py b/flax/linen/activation.py
@@ -69,6 +69,6 @@ def __call__(self, inputs: Array) -> Array:
     """
     negative_slope = self.param(
       'negative_slope',
-      lambda k: jnp.asarray(self.negative_slope_init)
+      lambda k: jnp.asarray(self.negative_slope_init, inputs.dtype)
     )
     return jnp.where(inputs >= 0, inputs, negative_slope * inputs)