https://github.com/JuliaReinforcementLearning/ReinforcementLearningCore.jl/blob/63f306d99a6db736a1755a5d1e26f2aa8e8822dc/src/extensions/Zygote.jl#L10 Or maybe make a PR in Flux instead?