diff --git a/CMakeLists.txt b/CMakeLists.txt
index 42f735ede964ffac3da68da5005e6b1e35b426b5..59ac688ce9b4bcf5ba80553ae6e8a598a53b844c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -8,12 +8,6 @@ project(aidge_learning
         DESCRIPTION "Functions and alogrithms to train models in the AIDGE framework" 
         LANGUAGES CXX)
 
-# Enable CUDA if backend_cuda is found
-find_package(aidge_backend_cuda)
-
-if(aidge_backend_cuda_FOUND)
-    enable_language(CUDA)
-endif()
 message(STATUS "Project name: ${CMAKE_PROJECT_NAME}")
 message(STATUS "Project version: ${version}")
 
@@ -26,6 +20,7 @@ option(PYBIND "python binding" OFF)
 option(WERROR "Warning as error" OFF)
 option(TEST "Enable tests" ON)
 option(COVERAGE "Enable coverage" OFF)
+option(ENABLE_ASAN "Enable ASan (AddressSanitizer) for runtime analysis of memory use (over/underflow, memory leak, ...)" OFF)
 
 ##############################################
 # Import utils CMakeLists
@@ -37,6 +32,7 @@ endif()
 
 ##############################################
 # FIND Dependencies
+
 if(NOT $ENV{AIDGE_INSTALL} STREQUAL "")
     set(CMAKE_INSTALL_PREFIX $ENV{AIDGE_INSTALL})
     list(APPEND CMAKE_PREFIX_PATH $ENV{AIDGE_INSTALL})
diff --git a/unit_tests/CMakeLists.txt b/unit_tests/CMakeLists.txt
index d33a87b486b8f14eab0cb8e99a9317ee14cbe647..39af9ec9b609b8b2c5685699ac014406ed64df72 100644
--- a/unit_tests/CMakeLists.txt
+++ b/unit_tests/CMakeLists.txt
@@ -1,55 +1,67 @@
-Include(FetchContent)
+include(FetchContent)
 
+# Fetch and make available Catch2 for unit testing
 FetchContent_Declare(
   Catch2
   GIT_REPOSITORY https://github.com/catchorg/Catch2.git
   GIT_TAG        v3.0.1 # or a later release
 )
-
 FetchContent_MakeAvailable(Catch2)
 
+# Gather all source files for the test executable
 file(GLOB_RECURSE src_files "*.cpp")
 set(tests_exe tests${module_name})
-add_executable(tests_exe ${src_files})
-
-set_target_properties(tests_exe PROPERTIES CUDA_SEPARABLE_COMPILATION ON)
 
+# Find required and optional dependencies
 find_package(aidge_backend_cpu REQUIRED)
 find_package(aidge_backend_cuda)
 
-if(NOT DEFINED CMAKE_CUDA_STANDARD)
-  set(CMAKE_CUDA_STANDARD 14)
-  set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+# Find CUDA Toolkit if CUDA backend is available
+if(aidge_backend_cuda_FOUND)
+  # Specify CUDA Toolkit without COMPONENTS to avoid missing imported targets
+  enable_language(CUDA)
+  find_package(CUDA REQUIRED)
+  find_package(CUDAToolkit REQUIRED)
+
+  if(NOT DEFINED CMAKE_CUDA_STANDARD)
+    set(CMAKE_CUDA_STANDARD 14)
+    set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+  endif()
 endif()
 
-target_link_libraries(tests_exe
+# Define the test executable
+add_executable(${tests_exe} ${src_files})
+
+# Include CUDA headers
+include_directories(${CUDA_INCLUDE_DIRS})
+# Link against the required aidge_backend_cpu library
+target_link_libraries(${tests_exe}
   PUBLIC ${module_name} _aidge_backend_cpu
   PRIVATE Catch2::Catch2WithMain
 )
 
+# Conditional CUDA linkage and definitions
 if(aidge_backend_cuda_FOUND)
+  # Enable CUDA language support and separable compilation for the target
   enable_language(CUDA)
-  find_package(CUDAToolkit REQUIRED)
-
-  target_link_libraries( _aidge_backend_cuda
-        INTERFACE
-        cudart  # Ensure CUDA dependencies are included
-        cublas  # Include other CUDA components if needed
-        cudnn
-  )
-
-  target_include_directories(_aidge_backend_cuda INTERFACE ${CUDAToolkit_INCLUDE_DIRS})
-
-  target_link_libraries(tests_exe
-        PUBLIC ${module_name} _aidge_backend_cuda
+  set_target_properties(${tests_exe} PROPERTIES CUDA_SEPARABLE_COMPILATION ON)
+  
+  target_include_directories(${tests_exe} PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
+  # Link manually specified CUDA libraries if the targets are not available
+  target_link_libraries(${tests_exe}
+    PUBLIC 
+      _aidge_backend_cuda
+      CUDA::cudart
+      CUDA::cublas
+      cudnn
   )
 
-  # Define a preprocessor macro if aidge_backend_cuda is linked
-  target_compile_definitions(tests_exe PUBLIC USE_AIDGE_BACKEND_CUDA)
+  # Define a preprocessor macro to indicate CUDA support
+  target_compile_definitions(${tests_exe} PUBLIC USE_AIDGE_BACKEND_CUDA)
 endif()
 
-
+# Add Catch2 test discovery
 list(APPEND CMAKE_MODULE_PATH ${catch2_SOURCE_DIR}/extras)
 include(CTest)
 include(Catch)
-catch_discover_tests(tests_exe)
+catch_discover_tests(${tests_exe})