fix: don't preserve structure in checks

LuxDL · Nov 14, 2024 · bcf5f64 · bcf5f64
1 parent dca6f85
commit bcf5f64
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/src/contrib/contrib.jl b/src/contrib/contrib.jl
@@ -6,7 +6,7 @@ using ChainRulesCore: ChainRulesCore
 using Compat: @compat
 using ConcreteStructs: @concrete
 using FastClosures: @closure
-using Functors: Functors, KeyPath, fmap_with_path, functor
+using Functors: Functors, KeyPath, fmap_with_path, fmapstructure_with_path, functor
 using Markdown: @doc_str
 using Optimisers: Optimisers
 using Random: AbstractRNG, Random

diff --git a/src/contrib/debug.jl b/src/contrib/debug.jl
@@ -91,7 +91,8 @@ function check_nan_and_throw(x, str::AbstractString, layer, location::KeyPath)
         return x
     end
 
-    return fmap_with_path(nan_check, x)
+    fmapstructure_with_path(nan_check, x)
+    return
 end
 
 function debug_layer_impl(layer, x, ps, st, location, error_check, _)