Update fast_gradient_method.py

fixes issue cleverhans-lab#1232
rekino · Jan 28, 2023 · 5c512e5 · 5c512e5
1 parent e5d00e5
commit 5c512e5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cleverhans/jax/attacks/fast_gradient_method.py b/cleverhans/jax/attacks/fast_gradient_method.py
@@ -1,6 +1,6 @@
 import jax.numpy as np
 from jax import grad, vmap
-from jax.experimental.stax import logsoftmax
+from jax.nn import log_softmax as logsoftmax
 
 from cleverhans.jax.utils import one_hot