diff --git a/test/instability_tests.jl b/test/instability_tests.jl
index 06f4169e..8fa11ef7 100644
--- a/test/instability_tests.jl
+++ b/test/instability_tests.jl
@@ -37,5 +37,6 @@ Test.@testset "Instability" begin
     JET.test_opt(
         ContinuousNormalizingFlows.loss,
         Base.typesof(icnf, ContinuousNormalizingFlows.TrainMode(), r, ps, st);
+        target_modules = [ContinuousNormalizingFlows],
     )
 end