Update shared_testsetup.jl: dont make other tests fail

JuliaGraphs · Aug 25, 2024 · 6227cd3 · 6227cd3
1 parent 4b32e2f
commit 6227cd3
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/GNNLux/test/shared_testsetup.jl b/GNNLux/test/shared_testsetup.jl
@@ -26,8 +26,13 @@ function test_lux_layer(rng::AbstractRNG, l, g::GNNGraph, x;
     st = LuxCore.initialstates(rng, l)
     @test LuxCore.parameterlength(l) == LuxCore.parameterlength(ps)
     @test LuxCore.statelength(l) == LuxCore.statelength(st)
-
-    y, st′ = l(g, x, edge_weight, ps, st)  
+
+    if edge_weight !== nothing
+        y, st′ = l(g, x, ps, st)
+    else          
+        y, st′ = l(g, x, edge_weight, ps, st)
+    end
+
     @test eltype(y) == eltype(x)
     if outputsize !== nothing
         @test LuxCore.outputsize(l) == outputsize
@@ -42,4 +47,4 @@ function test_lux_layer(rng::AbstractRNG, l, g::GNNGraph, x;
     test_gradients(loss, x, ps; atol, rtol, skip_backends=[AutoReverseDiff(), AutoTracker(), AutoForwardDiff(), AutoEnzyme()])
 end
 
-end
+end