fix typo in test

google · Sep 23, 2022 · 03abcc7 · 03abcc7
1 parent e76aa77
commit 03abcc7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/nn_test.py b/tests/nn_test.py
@@ -227,7 +227,7 @@ def fwd():
       a = jnp.array(1., 'float32')
 
       def f(hx, _):
-        hx = jax.nn.relu(hx + a)
+        hx = sigmoid(hx + a)
         return hx, None
 
       hx = jnp.array(0., 'float32')