fix params

This commit is contained in:
Iblis Lin 2018-02-08 12:55:16 +00:00 committed by Mike J Innes
parent 4511936a87
commit f7fdfbe3a9

View File

@ -33,9 +33,11 @@ function prefor(f, x; seen = OSet())
return return
end end
using Flux.Tracker: istracked
function params(m) function params(m)
ps = [] ps = []
prefor(p -> p isa TrackedArray && push!(ps, p), m) prefor(p -> istracked(p) && push!(ps, p), m)
return ps return ps
end end