# Help needed with Zygote Flux and custom adjoints on GPU

**URL:** <https://discourse.julialang.org/t/help-needed-with-zygote-flux-and-custom-adjoints-on-gpu/74326>\
**Category:** Machine Learning\
**Tags:** question, gpu, cuda, zygote, adjoint\
**Created:** [January 10, 2022, 11:21am UTC](https://discourse.julialang.org/t/help-needed-with-zygote-flux-and-custom-adjoints-on-gpu/74326 "2022-01-10T11:21:02Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![TsurHerman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tsurherman/32/1234_2.png) [@TsurHerman](https://discourse.julialang.org/u/TsurHerman)\
**Post date:** [January 10, 2022, 11:21am UTC](https://discourse.julialang.org/t/help-needed-with-zygote-flux-and-custom-adjoints-on-gpu/74326/1 "2022-01-10T11:21:02Z")

</div>

I am incorporating some simple non-standard elements in a larger DNN.

To achieve this I had to manually define some reverse rules through ChainRulesCore syntax

`ChainRulesCore.rrule(::typeof(f),args...)`

Which works great on the cpu, but somehow is not portable to the gpu.

I narrowed it to a MWE that fails the GPU compilation of the pullback function

```julia

using Zygote
using CUDA

f(A) = map(A) do a
    a
end

A = rand(Float32,128,128)
gA = gpu(A)
f(A) # works
f(gA) # works

Δ,pb = pullback(f,A) #works

Δ,pb = pullback(f,gA) #fails

```

Any insights?

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [January 11, 2022, 3:27am UTC](https://discourse.julialang.org/t/help-needed-with-zygote-flux-and-custom-adjoints-on-gpu/74326/2 "2022-01-11T03:27:30Z")

</div>

I don’t believe `map` works in reverse on GPU with Zygote. Workarounds are to define an rrule that encapsulates the map, or use broadcasting instead.