don't use ReverseWithPrimal in gradient

FluxML · Nov 11, 2024 · a9424ab · a9424ab
1 parent 2636454
commit a9424ab
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ext/FluxEnzymeExt/FluxEnzymeExt.jl b/ext/FluxEnzymeExt/FluxEnzymeExt.jl
@@ -20,7 +20,7 @@ function Flux._enzyme_gradient(f, args::Union{Const, Duplicated}...; zero::Bool=
     zero && x isa Duplicated && _make_zero!(x.dval)
     _check_mutable(x)
   end
-  Enzyme.autodiff(ReverseWithPrimal, f, Active, args...)
+  Enzyme.autodiff(Reverse, f, Active, args...)
   map(_grad_or_nothing, args)
 end