jax.nn.functions has moved to jax.nn

matpalm · Jan 31, 2021 · dd40127 · dd40127
1 parent 973c42f
commit dd40127
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models.py b/models.py
@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 from jax import random, lax, vmap
 from jax.nn.initializers import glorot_normal, he_normal
-from jax.nn.functions import gelu
+from jax.nn import gelu
 from functools import partial
 import objax
 from objax.variable import TrainVar, StateVar