Merge pull request #679 from SciML/Vaibhavdixit02-patch-4

Update OptimizationZygoteExt.jl
SciML · Jan 11, 2024 · d6bea20 · d6bea20
2 parents f2001e9 + c7347d2
commit d6bea20
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ext/OptimizationZygoteExt.jl b/ext/OptimizationZygoteExt.jl
@@ -13,7 +13,7 @@ function Optimization.instantiate_function(f, x, adtype::AutoZygote, p,
         grad = function (res, θ, args...)
             val = Zygote.gradient(x -> _f(x, args...), θ)[1]
             if val === nothing
-                res .= 0
+                res .= zero(typeof(θ))
             else
                 res .= val
             end
@@ -90,7 +90,7 @@ function Optimization.instantiate_function(f, cache::Optimization.ReInitCache,
         grad = function (res, θ, args...)
             val = Zygote.gradient(x -> _f(x, args...), θ)[1]
             if val === nothing
-                res .= 0
+                res .= zero(typeof(θ))
             else
                 res .= val
             end