fix: reset inside module

LuxDL · Jul 13, 2024 · d8d69e2 · d8d69e2
1 parent 70e4f89
commit d8d69e2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxTestUtils"
 uuid = "ac9de150-d08f-4546-94fb-7472b5760531"
 authors = ["Avik Pal <[email protected]>"]
-version = "0.1.18"
+version = "0.1.19"
 
 [deps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"

diff --git a/src/LuxTestUtils.jl b/src/LuxTestUtils.jl
@@ -6,7 +6,11 @@ using ForwardDiff, ReverseDiff, Tracker, Zygote, FiniteDifferences
 const JET_TARGET_MODULES = Ref{Union{Nothing, Vector{String}}}(nothing)
 
 function __init__()
-    return JET_TARGET_MODULES[] = @load_preference("target_modules", nothing)
+    if @has_preference("target_modules")
+        prefs = @load_preference("target_modules")
+        @info "JET_TARGET_MODULES set to $prefs from preferences"
+        JET_TARGET_MODULES[] = prefs
+    end
 end
 
 function jet_target_modules!(list::Vector{String})
@@ -90,6 +94,7 @@ macro jet(expr, args...)
         if !target_modules_set && JET_TARGET_MODULES[] !== nothing
             target_modules = getproperty.(
                 (__module__,), Tuple(Symbol.(JET_TARGET_MODULES[])))
+            @show target_modules
             push!(all_args, :(target_modules = $target_modules))
         end