Merge pull request #350 from FluxML/patch

Fix failure as deleted
FluxML · Jan 8, 2023 · 63265df · 63265df
2 parents 312325a + 8b68946
commit 63265df
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/operation.jl b/src/operation.jl
@@ -43,7 +43,7 @@ function indexed_softmax(x::AbstractArray, xs, N; dims=1)
 end
 
 function ∇indexed_softmax(dy::AbstractArray{T}, y::AbstractArray{S}, xs, N; dims=1) where {T,S}
-    dx = if NNlib.within_grad()
+    dx = if NNlib.within_gradient(y)
         tmp = dy .* y
         for i in 1:N
             idx = ntuple(j -> (j == dims) ? (xs .== i) : Colon(), ndims(y))