From 6227cd368f02500790b7907846584d5b6b76fc31 Mon Sep 17 00:00:00 2001
From: Rishabh <59335537+rbSparky@users.noreply.github.com>
Date: Mon, 26 Aug 2024 03:39:37 +0530
Subject: [PATCH] Update shared_testsetup.jl: dont make other tests fail

---
 GNNLux/test/shared_testsetup.jl | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/GNNLux/test/shared_testsetup.jl b/GNNLux/test/shared_testsetup.jl
index 332b578c6..bcd243df3 100644
--- 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
\ No newline at end of file
+end