Skip to content

Commit 7957bd6

Browse files
Skip parameter updates when gradients contain NaN or Inf
- Add check before Optimisers.update to detect NaN/Inf in gradients - Skip update but still increment iteration counter when detected - Add warning message (maxlog=10) to inform users - Fixes issue where NaN gradients corrupt all subsequent updates Addresses: https://discourse.julialang.org/t/how-to-ignore-minibatches-with-nan-gradients-optimizing-a-hybrid-lux-model-using-optimization-jl/132615
1 parent 846e6c1 commit 7957bd6

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

lib/OptimizationOptimisers/src/OptimizationOptimisers.jl

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,12 @@ function SciMLBase.__solve(cache::OptimizationCache{O}) where {O <: AbstractRule
129129
break
130130
end
131131
end
132-
state, θ = Optimisers.update(state, θ, G)
132+
# Skip update if gradient contains NaN or Inf values
133+
if !any(x -> isnan(x) || isinf(x), G)
134+
state, θ = Optimisers.update(state, θ, G)
135+
else
136+
@warn "Skipping parameter update due to NaN or Inf in gradients at iteration $iterations" maxlog=10
137+
end
133138
end
134139
cache.progress && @logmsg(LogLevel(-1), "Optimization",
135140
_id=progress_id, message="Done", progress=1.0)

0 commit comments

Comments
 (0)