smaller computational graph

FluxML · Aug 29, 2024 · a816377 · a816377
1 parent 2febf9f
commit a816377
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/lib/array.jl b/src/lib/array.jl
@@ -144,7 +144,7 @@ logabsdet(xs::TrackedArray) = track(logabsdet, xs)
 @grad logabsdet(xs) = logabsdet(data(xs)), Δ -> (Δ[1] * transpose(inv(xs)),)
 
 Base.repeat(xs::TrackedArray; kw...) = track(repeat, xs; kw...)
-Base.zero(x::Tracker.TrackedArray) = zero.(x)
+Base.zero(x::Tracker.TrackedArray) = TrackedArray(zero(x.data))
 
 @grad function repeat(xs; inner=ntuple(x->1, ndims(xs)), outer=ntuple(x->1, ndims(xs)))
   repeat(data(xs), inner = inner, outer = outer), function (Δ)