diff --git a/GNNGraphs/Project.toml b/GNNGraphs/Project.toml
index bd3b83340..27a1fa6df 100644
--- a/GNNGraphs/Project.toml
+++ b/GNNGraphs/Project.toml
@@ -31,7 +31,7 @@ GNNGraphsSimpleWeightedGraphsExt = "SimpleWeightedGraphs"
 Adapt = "4"
 CUDA = "5"
 ChainRulesCore = "1"
-Functors = "0.4.1"
+Functors = "0.4.1, 0.5"
 Graphs = "1.4"
 KrylovKit = "0.8"
 LinearAlgebra = "1"
diff --git a/GNNLux/LICENSE b/GNNLux/LICENSE
index e7410390d..e58d7f0be 100644
--- a/GNNLux/LICENSE
+++ b/GNNLux/LICENSE
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2021 Carlo Lucibello <carlo.lucibello@gmail.com> and contributors
+Copyright (c) 2024 Carlo Lucibello <carlo.lucibello@gmail.com> and contributors
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/GNNLux/Project.toml b/GNNLux/Project.toml
index 54944c740..46593085e 100644
--- a/GNNLux/Project.toml
+++ b/GNNLux/Project.toml
@@ -17,22 +17,12 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 ConcreteStructs = "0.2.3"
-Lux = "1.0"
-LuxCore = "1.0"
+GNNGraphs = "1.3"
+GNNlib = "0.2.3"
+Lux = "1"
+LuxCore = "1"
 NNlib = "0.9.21"
 Reexport = "1.2"
 Static = "1.1"
+Statistics = "1"
 julia = "1.10"
-
-[extras]
-ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
-Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
-LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
-MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
-ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
-StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-
-[targets]
-test = ["Test", "MLDataDevices", "ComponentArrays", "Functors", "LuxTestUtils", "ReTestItems", "StableRNGs", "Zygote"]
diff --git a/GNNLux/test/Project.toml b/GNNLux/test/Project.toml
new file mode 100644
index 000000000..d651808bf
--- /dev/null
+++ b/GNNLux/test/Project.toml
@@ -0,0 +1,22 @@
+[deps]
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
+GNNLux = "e8545f4d-a905-48ac-a8c4-ca114b98986d"
+GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
+LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
+MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
+TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
diff --git a/GNNLux/test/layers/basic_tests.jl b/GNNLux/test/layers/basic.jl
similarity index 86%
rename from GNNLux/test/layers/basic_tests.jl
rename to GNNLux/test/layers/basic.jl
index f4cabad69..0064d12e6 100644
--- a/GNNLux/test/layers/basic_tests.jl
+++ b/GNNLux/test/layers/basic.jl
@@ -1,4 +1,6 @@
-@testitem "layers/basic" setup=[SharedTestSetup] begin
+@testitem "layers/basic" setup=[TestModuleLux] begin
+    using .TestModuleLux
+    
     rng = StableRNG(17)
     g = rand_graph(rng, 10, 40)
     x = randn(rng, Float32, 3, 10)        
diff --git a/GNNLux/test/layers/conv_tests.jl b/GNNLux/test/layers/conv.jl
similarity index 98%
rename from GNNLux/test/layers/conv_tests.jl
rename to GNNLux/test/layers/conv.jl
index 4f871b64e..dd30795ee 100644
--- a/GNNLux/test/layers/conv_tests.jl
+++ b/GNNLux/test/layers/conv.jl
@@ -1,4 +1,6 @@
-@testitem "layers/conv" setup=[SharedTestSetup] begin
+@testitem "layers/conv" setup=[TestModuleLux] begin
+    using .TestModuleLux
+    
     rng = StableRNG(1234)
     g = rand_graph(rng, 10, 40)
     in_dims = 3
diff --git a/GNNLux/test/layers/temporalconv_test.jl b/GNNLux/test/layers/temporalconv.jl
similarity index 96%
rename from GNNLux/test/layers/temporalconv_test.jl
rename to GNNLux/test/layers/temporalconv.jl
index ec670b6bc..8de93efa7 100644
--- a/GNNLux/test/layers/temporalconv_test.jl
+++ b/GNNLux/test/layers/temporalconv.jl
@@ -1,4 +1,5 @@
-@testitem "layers/temporalconv" setup=[SharedTestSetup] begin
+@testitem "layers/temporalconv" setup=[TestModuleLux] begin
+    using .TestModuleLux
     using LuxTestUtils: test_gradients, AutoReverseDiff, AutoTracker, AutoForwardDiff, AutoEnzyme
 
     rng = StableRNG(1234)
diff --git a/GNNLux/test/runtests.jl b/GNNLux/test/runtests.jl
index 8bd6c7ff1..91f74105f 100644
--- a/GNNLux/test/runtests.jl
+++ b/GNNLux/test/runtests.jl
@@ -1,10 +1,31 @@
-using Test
-using Lux
-using GNNLux
-using Random, Statistics
+using TestItemRunner
 
-using ReTestItems
-# using Pkg, Preferences, Test
-# using InteractiveUtils, Hwloc
+## See https://www.julia-vscode.org/docs/stable/userguide/testitems/
+## for how to run the tests within VS Code.
+## See test_module.jl for the test infrastructure.
 
-runtests(GNNLux)
+## Uncomment below and in test_module.jl to change the default test settings
+# ENV["GNN_TEST_CPU"] = "false"
+# ENV["GNN_TEST_CUDA"] = "true"
+# ENV["GNN_TEST_AMDGPU"] = "true"
+# ENV["GNN_TEST_Metal"] = "true"
+
+# The only available tag at the moment is :gpu
+# Tests not tagged with :gpu are considered to be CPU tests
+# Tests tagged with :gpu should run on all GPU backends
+
+# TODO add gpu tests. Not urgent since almost the whole path for layers is tested 
+# in GraphNeuralNetworks.jl
+
+if get(ENV, "GNN_TEST_CPU", "true") == "true"
+    @run_package_tests filter = ti -> :gpu ∉ ti.tags
+end
+if get(ENV, "GNN_TEST_CUDA", "false") == "true"
+    @run_package_tests filter = ti -> :gpu ∈ ti.tags
+end
+if get(ENV, "GNN_TEST_AMDGPU", "false") == "true"
+    @run_package_tests filter = ti -> :gpu ∈ ti.tags
+end
+if get(ENV, "GNN_TEST_Metal", "false") == "true"
+    @run_package_tests filter = ti -> :gpu ∈ ti.tags
+end
diff --git a/GNNLux/test/shared_testsetup.jl b/GNNLux/test/test_module.jl
similarity index 65%
rename from GNNLux/test/shared_testsetup.jl
rename to GNNLux/test/test_module.jl
index bf2bdbbf2..83f3e7062 100644
--- a/GNNLux/test/shared_testsetup.jl
+++ b/GNNLux/test/test_module.jl
@@ -1,6 +1,30 @@
-@testsetup module SharedTestSetup
+@testmodule TestModuleLux begin
 
-import Reexport: @reexport
+using Pkg
+
+## Uncomment below to change the default test settings
+# ENV["GNN_TEST_CUDA"] = "true"
+# ENV["GNN_TEST_AMDGPU"] = "true"
+# ENV["GNN_TEST_Metal"] = "true"
+
+to_test(backend) = get(ENV, "GNN_TEST_$(backend)", "false") == "true"
+has_dependecies(pkgs) = all(pkg -> haskey(Pkg.project().dependencies, pkg), pkgs)
+deps_dict = Dict(:CUDA => ["CUDA", "cuDNN"], :AMDGPU => ["AMDGPU"], :Metal => ["Metal"])
+
+for (backend, deps) in deps_dict
+    if to_test(backend)
+        if !has_dependecies(deps)
+            Pkg.add(deps)
+        end
+        @eval using $backend
+        if backend == :CUDA
+            @eval using cuDNN
+        end
+        @eval $backend.allowscalar(false)
+    end
+end
+
+using Reexport: @reexport
 
 @reexport using Test
 @reexport using GNNLux