fix: missing []

LuxDL · Jul 12, 2024 · 70e4f89 · 70e4f89 · avik-pal · Jul 12, 2024
1 parent d24bae7
commit 70e4f89
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 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.17"
+version = "0.1.18"
 
 [deps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"

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