patch

SciML · Jun 30, 2022 · 2b563a5 · 2b563a5
1 parent e32422d
commit 2b563a5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "DiffEqFlux"
 uuid = "aae7a2af-3d4f-5e19-a356-7da93b79d9d0"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "1.51.0"
+version = "1.51.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/fast_layers.jl b/src/fast_layers.jl
@@ -120,7 +120,6 @@ ZygoteRules.@adjoint function (f::FastDense)(x::Number,p)
           tmp = typeof(bbar) <: AbstractVector ?
                            vec(vcat(vec(Wbar),bbar)) :
                            vec(vcat(vec(Wbar),sum(bbar,dims=2)))
-          ifgpufree(bbar)
           tmp
       else
           vec(Wbar)
@@ -203,7 +202,7 @@ ZygoteRules.@adjoint function (f::FastDense)(x::AbstractVector,p)
           tmp = typeof(bbar) <: AbstractVector ?
                            vec(vcat(vec(Wbar),bbar)) :
                            vec(vcat(vec(Wbar),sum(bbar,dims=2)))
-          ifgpufree(bbar)
+          !(typeof(f.σ) <: typeof(identity)) && ifgpufree(bbar)
           tmp
       else
           vec(Wbar)