Tracker.jl icon indicating copy to clipboard operation
Tracker.jl copied to clipboard

Gradients dropped by `adapt`

Open mcabbott opened this issue 1 year ago • 1 comments

Moving y to "gpu" inside loss causes its gradient to be lost:

julia> using Tracker, JLArrays

julia> JLArrays.allowscalar(false)

julia> Tracker.withgradient((x,y) -> sum(x[1:2] + jl(y))^2, jl([1,2,3.0]), [4,5.0])
(val = 144.0, grad = ([24.0, 24.0, 0.0], [0.0, 0.0])) 

julia> ans.grad[1] isa JLArray
true

unlike Zygote:

julia> Zygote.withgradient((x,y) -> sum(x[1:2] + jl(y))^2, jl([1,2,3.0]), [4,5.0])
(val = 144.0, grad = ([24.0, 24.0, 0.0], [24.0, 24.0]))

mcabbott avatar Aug 22 '22 00:08 mcabbott