Repository: lhnguyen102/cuTAGI (language: python)

Issue:
Unit test using CNN with batch normalization
Added CNN batch. norm unit test and solved some errors in the main script.


Discussion / hints:

Fix patch:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index b082ad692..5b96cac42 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -159,9 +159,11 @@ set(GPU_SOURCES
   src/tagi_network.cu
   src/gpu_debug_utils.cpp
   test/test_gpu.cu
+  test/autoencoder/test_autoencoder_gpu.cu
   test/cnn/test_cnn_gpu.cu
   test/cnn_batch_norm/test_cnn_batch_norm_gpu.cu
   test/test_classification.cu
+  test/test_autoencoder.cu
 )
 
 # Output binary folder der for different mode
diff --git a/main.cpp b/main.cpp
index b3535f50d..3e4e29ed2 100644
--- a/main.cpp
+++ b/main.cpp
@@ -43,7 +43,8 @@ int main(int argc, char *argv[]) {
     if (user_input_file.compare("test") == 0) {
         // test_lstm_cpu();
         bool compute_gpu_tests = false;
-        test_cpu(user_input_options, compute_gpu_tests);
+        auto start = std::chrono::steady_clock::now();
+        test_cpu(user_input_options, compute_gpu_tests, start);
         std::cout << std::endl;
         std::cout << "Unable to perform test on GPU: CUDA device unavailable."
                   << std::endl;
diff --git a/main.cu b/main.cu
index df6df0d4b..4a1871987 100644
--- a/main.cu
+++ b/main.cu
@@ -48,11 +48,12 @@ int main(int argc, char* argv[]) {
     if (user_input_file.compare("test") == 0) {
         // auto a = test_lstm_cpu();
         bool compute_gpu_tests = is_cuda_available();
+        auto start = std::chrono::steady_clock::now();
         int num_tests_passed_cpu =
-            test_cpu(user_input_options, compute_gpu_tests);
+            test_cpu(user_input_options, compute_gpu_tests, start);
         // If cuda Available and output reinicialization not aborted
         if (compute_gpu_tests && num_tests_passed_cpu >= 0)
-            test_gpu(user_input_options, num_tests_passed_cpu);
+            test_gpu(user_input_options, num_tests_passed_cpu, start);
         else if (!compute_gpu_tests) {
             std::cout << std::endl;
             std::cout


Test patch:
diff --git a/test/act_func/test_act_func_cpu.cpp b/test/act_func/test_act_func_cpu.cpp
index 996c95139..f94f42336 100644
--- a/test/act_func/test_act_func_cpu.cpp
+++ b/test/act_func/test_act_func_cpu.cpp
@@ -32,14 +32,8 @@ bool test_act_func_cpu(bool recompute_outputs, std::string date,
     path.curr_path = get_current_dir();
     std::string data_path = path.curr_path + "/test/data/" + data;
     std::string data_dir = path.curr_path + "/test/" + arch + "/data/";
-    std::string init_param_path_w =
-        data_dir + date + "_" + data + "_init_param_weights_w.csv";
-    std::string init_param_path_w_sc =
-        data_dir + date + "_" + data + "_init_param_weights_w_sc.csv";
-    std::string init_param_path_b =
-        data_dir + date + "_" + data + "_init_param_bias_b.csv";
-    std::string init_param_path_b_sc =
-        data_dir + date + "_" + data + "_init_param_bias_b_sc.csv";
+
+    TestSavingPaths test_saving_paths(path.curr_path, arch, data, date);
 
     // Iterate over all activation functions
     for (int i = 0; i < ACTIVATIONS.size(); i++) {
@@ -66,19 +60,7 @@ bool test_act_func_cpu(bool recompute_outputs, std::string date,
         Dataloader test_db =
             test_data(data, tagi_net, data_path, train_db, NORMALIZE);
 
-        std::vector<std::vector<float> *> weights;
-        weights.push_back(&tagi_net.theta.mw);
-        weights.push_back(&tagi_net.theta.Sw);
-        std::vector<std::vector<float> *> weights_sc;
-        weights_sc.push_back(&tagi_net.theta.mw_sc);
-        weights_sc.push_back(&tagi_net.theta.Sw_sc);
-
-        std::vector<std::vector<float> *> bias;
-        bias.push_back(&tagi_net.theta.mb);
-        bias.push_back(&tagi_net.theta.Sb);
-        std::vector<std::vector<float> *> bias_sc;
-        bias_sc.push_back(&tagi_net.theta.mb_sc);
-        bias_sc.push_back(&tagi_net.theta.Sb_sc);
+        TestParamAndStates params_and_states(tagi_net);
 
         // If we want to test but no data is available, we throw an error
         if (!recompute_outputs && !directory_exists(data_dir)) {
@@ -92,49 +74,40 @@ bool test_act_func_cpu(bool recompute_outputs, std::string date,
                           << std::endl;
                 return false;
             }
-            write_vector_to_csv(init_param_path_w, "mw,Sw", weights);
-            write_vector_to_csv(init_param_path_w_sc, "mw_sc,Sw_sc",
-                                weights_sc);
-
-            write_vector_to_csv(init_param_path_b, "mb,Sb", bias);
-            write_vector_to_csv(init_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+            params_and_states.write_params(test_saving_paths, true);
         }
 
         // Read the initial parameters (see tes_utils.cpp for more details)
-        read_vector_from_csv(init_param_path_w, weights);
-        read_vector_from_csv(init_param_path_w_sc, weights_sc);
-        read_vector_from_csv(init_param_path_b, bias);
-        read_vector_from_csv(init_param_path_b_sc, bias_sc);
+        params_and_states.read_params(test_saving_paths, true);
 
         // Train the network
         forward_pass(tagi_net, train_db);
 
-        std::vector<std::vector<float> *> forward_states;
-        forward_states.push_back(&tagi_net.state.mz);
-        forward_states.push_back(&tagi_net.state.Sz);
-        forward_states.push_back(&tagi_net.state.ma);
-        forward_states.push_back(&tagi_net.state.Sa);
-        forward_states.push_back(&tagi_net.state.J);
+        add_forward_states(params_and_states.forward_states, tagi_net);
 
         if (recompute_outputs) {
             // RESET OUPUTS
 
             // Write the forward hidden states
             write_vector_to_csv(forward_states_path, "mz,Sz,ma,Sa,J",
-                                forward_states);
+                                params_and_states.forward_states);
 
         } else {
             // PERFORM TESTS
 
+            TestParamAndStates params_and_states_reference(tagi_net);
+
             // Read the saved forward hidden states reference
-            std::vector<std::vector<float> *> ref_forward_states;
             for (int i = 0; i < 5; i++)
-                ref_forward_states.push_back(new std::vector<float>());
+                params_and_states_reference.forward_states.push_back(
+                    new std::vector<float>());
 
-            read_vector_from_csv(forward_states_path, ref_forward_states);
+            read_vector_from_csv(forward_states_path,
+                                 params_and_states_reference.forward_states);
 
             // Compare the saved forward hidden states with the ones we got
-            if (!compare_vectors(ref_forward_states, forward_states, data,
+            if (!compare_vectors(params_and_states_reference.forward_states,
+                                 params_and_states.forward_states, data,
                                  "fnn forward hidden states")) {
                 std::cout << "\033[1;31mTest for " + ACTIVATIONS_NAMES[i] +
                                  " " + "activation function has FAILED in " +
diff --git a/test/autoencoder/test_autoencoder_gpu.cu b/test/autoencoder/test_autoencoder_gpu.cu
new file mode 100644
index 000000000..8837154fb
--- /dev/null
+++ b/test/autoencoder/test_autoencoder_gpu.cu
@@ -0,0 +1,325 @@
+///////////////////////////////////////////////////////////////////////////////
+// File:         test_autoencoder_gpu.cu
+// Description:  Script to test the autoencoder GPU implementation of cuTAGI
+// Authors:      Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet
+// Created:      April 11, 2023
+// Updated:      April 16, 2023
+// Contact:      miquelflorensa11@gmail.com & luongha.nguyen@gmail.com &
+//               james.goulet@polymtl.ca
+// Copyright (c) 2023 Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet.
+// Some rights reserved.
+///////////////////////////////////////////////////////////////////////////////
+
+#include "test_autoencoder_gpu.cuh"
+
+// Specify network properties for the encoder
+const std::vector<int> LAYERS_E = {2, 2, 6, 4, 2, 6, 4, 1, 1};
+const std::vector<int> NODES_E = {784, 0, 0, 0, 0, 0, 0, 100, 10};
+const std::vector<int> KERNELS_E = {3, 1, 3, 3, 1, 3, 1, 1, 1};
+const std::vector<int> STRIDES_E = {1, 0, 2, 1, 0, 2, 0, 0, 0};
+const std::vector<int> WIDTHS_E = {28, 0, 0, 0, 0, 0, 0, 0, 0};
+const std::vector<int> HEIGHTS_E = {28, 0, 0, 0, 0, 0, 0, 0, 0};
+const std::vector<int> FILTERS_E = {1, 4, 4, 4, 8, 8, 8, 1, 1};
+const std::vector<int> PADS_E = {1, 0, 1, 1, 0, 1, 0, 0, 0};
+const std::vector<int> PAD_TYPES_E = {1, 0, 2, 1, 0, 2, 0, 0, 0};
+const std::vector<int> ACTIVATIONS_E = {0, 4, 0, 0, 4, 0, 0, 4, 0};
+
+// Specify network properties for the decoder
+const std::vector<int> LAYERS_D = {1, 1, 21, 21, 21};
+const std::vector<int> NODES_D = {10, 1568, 0, 0, 784};
+const std::vector<int> KERNELS_D = {1, 3, 3, 3, 1};
+const std::vector<int> STRIDES_D = {0, 2, 2, 1, 0};
+const std::vector<int> WIDTHS_D = {0, 7, 0, 0, 0};
+const std::vector<int> HEIGHTS_D = {0, 7, 0, 0, 0};
+const std::vector<int> FILTERS_D = {1, 8, 8, 4, 1};
+const std::vector<int> PADS_D = {0, 1, 1, 1, 0};
+const std::vector<int> PAD_TYPES_D = {0, 2, 2, 1, 0};
+const std::vector<int> ACTIVATIONS_D = {0, 4, 4, 4, 0};
+
+const int BATCH_SIZE = 2;
+const int SIGMA_V = 8;
+const int SIGMA_V_MIN = 2;
+const int DECAT_FACTOR_SIGMA_V = 0.95;
+const int NUM_CLASSES = 10;
+const std::vector<float> MU = {0.1309};
+const std::vector<float> SIGMA = {1.0};
+const std::string INIT_METHOD = "He";
+
+bool test_autoencoder_gpu(bool recompute_outputs, std::string date,
+                          std::string arch, std::string data) {
+    // Encoder
+    Network net_prop_e;
+
+    net_prop_e.layers = LAYERS_E;
+    net_prop_e.nodes = NODES_E;
+    net_prop_e.kernels = KERNELS_E;
+    net_prop_e.strides = STRIDES_E;
+    net_prop_e.widths = WIDTHS_E;
+    net_prop_e.heights = HEIGHTS_E;
+    net_prop_e.filters = FILTERS_E;
+    net_prop_e.pads = PADS_E;
+    net_prop_e.pad_types = PAD_TYPES_E;
+    net_prop_e.activations = ACTIVATIONS_E;
+    net_prop_e.batch_size = BATCH_SIZE;
+    net_prop_e.init_method = INIT_METHOD;
+
+    std::string device = "cuda";
+    net_prop_e.device = device;
+
+    TagiNetwork net_e(net_prop_e);
+    net_e.prop.is_output_ud = false;
+
+    // Decoder
+    Network net_prop_d;
+    net_prop_d.layers = LAYERS_D;
+    net_prop_d.nodes = NODES_D;
+    net_prop_d.kernels = KERNELS_D;
+    net_prop_d.strides = STRIDES_D;
+    net_prop_d.widths = WIDTHS_D;
+    net_prop_d.heights = HEIGHTS_D;
+    net_prop_d.filters = FILTERS_D;
+    net_prop_d.pads = PADS_D;
+    net_prop_d.pad_types = PAD_TYPES_D;
+    net_prop_d.activations = ACTIVATIONS_D;
+    net_prop_d.batch_size = BATCH_SIZE;
+    net_prop_d.sigma_v = SIGMA_V;
+    net_prop_d.sigma_v_min = SIGMA_V_MIN;
+    net_prop_d.decay_factor_sigma_v = DECAT_FACTOR_SIGMA_V;
+    net_prop_d.init_method = INIT_METHOD;
+
+    net_prop_d.device = device;
+
+    TagiNetwork net_d(net_prop_d);
+    net_d.prop.last_backward_layer = 0;
+
+    SavePath path;
+    path.curr_path = get_current_dir();
+
+    TestSavingPaths test_saving_paths_encoder(path.curr_path, arch, data, date,
+                                              true);
+    TestSavingPaths test_saving_paths_decoder(path.curr_path, arch, data, date,
+                                              false, true);
+
+    std::string data_path = path.curr_path + "/test/data/" + data;
+    std::string data_dir = path.curr_path + "/test/" + arch + "/data/";
+
+    // Train data
+    ImageData imdb =
+        image_dataloader(data, data_path, MU, SIGMA, NUM_CLASSES, net_e.prop);
+
+    TestParamAndStates params_and_states_encoder(net_e);
+    TestParamAndStates params_and_states_decoder(net_d);
+
+    // If we want to test but no data is available, we throw an error
+    if (!recompute_outputs && !directory_exists(data_dir)) {
+        throw std::runtime_error(
+            "Tested data are not available. Please regenerate or provide the "
+            "tested data");
+    } else if (recompute_outputs) {
+        // Cheks if the data directory exists and if not, creates it
+        if (!create_directory_if_not_exists(data_dir)) {
+            std::cout << "Error: could not create data directory" << std::endl;
+            return false;
+        }
+
+        params_and_states_encoder.write_params(test_saving_paths_encoder, true);
+        params_and_states_decoder.write_params(test_saving_paths_decoder, true);
+    }
+
+    // Read the initial parameters (see test_utils.cpp for more details)
+    params_and_states_encoder.read_params(test_saving_paths_encoder, true);
+    params_and_states_decoder.read_params(test_saving_paths_decoder, true);
+
+    net_e.theta_gpu.copy_host_to_device();
+    net_d.theta_gpu.copy_host_to_device();
+
+    // Autoencoder
+    train_autoencoder(net_e, net_d, imdb, NUM_CLASSES);
+
+    // Transfer data to CPU
+    net_e.theta_gpu.copy_device_to_host();
+    net_d.theta_gpu.copy_device_to_host();
+    net_e.d_state_gpu.copy_device_to_host();
+    net_d.d_state_gpu.copy_device_to_host();
+
+    // Recover forward and backward hidden states after running the autoencoder
+    add_forward_states(params_and_states_encoder.forward_states, net_e);
+    add_forward_states(params_and_states_decoder.forward_states, net_d);
+
+    // Recover forward and backward hidden states after running the autoencoder
+
+    std::vector<std::vector<float>> backward_states_e;
+    std::vector<std::vector<float>> backward_states_d;
+    std::string backward_states_header_e = "";
+    std::string backward_states_header_d = "";
+
+    add_backward_states(backward_states_e, backward_states_header_e, net_e,
+                        net_prop_e.layers.size());
+    add_backward_states(backward_states_d, backward_states_header_d, net_d,
+                        net_prop_d.layers.size());
+
+    for (int i = 0; i < backward_states_e.size(); i++)
+        params_and_states_encoder.backward_states.push_back(
+            &backward_states_e[i]);
+
+    for (int i = 0; i < backward_states_d.size(); i++)
+        params_and_states_decoder.backward_states.push_back(
+            &backward_states_d[i]);
+
+    if (recompute_outputs) {
+        // RESET OUPUTS
+
+        // Write the parameters and hidden states
+        params_and_states_encoder.write_params(test_saving_paths_encoder,
+                                               false);
+        params_and_states_decoder.write_params(test_saving_paths_decoder,
+                                               false);
+
+        // Write the forward hidden states
+        write_vector_to_csv(test_saving_paths_encoder.forward_states_path,
+                            "mz,Sz,ma,Sa,J",
+                            params_and_states_encoder.forward_states);
+        write_vector_to_csv(test_saving_paths_decoder.forward_states_path,
+                            "mz,Sz,ma,Sa,J",
+                            params_and_states_decoder.forward_states);
+
+        // Write the backward hidden states
+        write_vector_to_csv(test_saving_paths_encoder.backward_states_path,
+                            backward_states_header_e,
+                            params_and_states_encoder.backward_states);
+        write_vector_to_csv(test_saving_paths_decoder.backward_states_path,
+                            backward_states_header_d,
+                            params_and_states_decoder.backward_states);
+
+    } else {
+        // PERFORM TESTS
+
+        // Read the saved reference parameters
+        TestParamAndStates params_and_states_encoder_reference(net_e);
+        TestParamAndStates params_and_states_decoder_reference(net_d);
+
+        params_and_states_encoder_reference.read_params(
+            test_saving_paths_encoder, false);
+        params_and_states_decoder_reference.read_params(
+            test_saving_paths_decoder, false);
+
+        net_e.theta_gpu.copy_host_to_device();
+        net_d.theta_gpu.copy_host_to_device();
+
+        // Compare optimal values with the ones we got
+        if (!compare_vectors(params_and_states_encoder_reference.weights,
+                             params_and_states_encoder.weights, data,
+                             "encoder weights") ||
+            !compare_vectors(params_and_states_encoder_reference.weights_sc,
+                             params_and_states_encoder.weights_sc, data,
+                             "encoder weights for residual network") ||
+            !compare_vectors(params_and_states_encoder_reference.bias,
+                             params_and_states_encoder.bias, data,
+                             "encoder bias") ||
+            !compare_vectors(params_and_states_encoder_reference.bias_sc,
+                             params_and_states_encoder.bias_sc, data,
+                             "encoder bias for residual network")) {
+            std::cout << "\033[1;31mTest for encoder PARAMS has FAILED in " +
+                             data + " data\033[0m\n"
+                      << std::endl;
+            return false;
+        }
+
+        if (!compare_vectors(params_and_states_decoder_reference.weights,
+                             params_and_states_decoder.weights, data,
+                             "decoder weights") ||
+            !compare_vectors(params_and_states_decoder_reference.weights_sc,
+                             params_and_states_decoder.weights_sc, data,
+                             "decoder weights for residual network") ||
+            !compare_vectors(params_and_states_decoder_reference.bias,
+                             params_and_states_decoder.bias, data,
+                             "decoder bias") ||
+            !compare_vectors(params_and_states_decoder_reference.bias_sc,
+                             params_and_states_decoder.bias_sc, data,
+                             "decoder bias for residual network")) {
+            std::cout << "\033[1;31mTest for decoder PARAMS has FAILED in " +
+                             data + " data\033[0m\n"
+                      << std::endl;
+            return false;
+        }
+
+        // Read the saved forward hidden states reference
+        for (int i = 0; i < 5; i++)
+            params_and_states_encoder_reference.forward_states.push_back(
+                new std::vector<float>());
+        for (int i = 0; i < 5; i++)
+            params_and_states_decoder_reference.forward_states.push_back(
+                new std::vector<float>());
+
+        read_vector_from_csv(
+            test_saving_paths_encoder.forward_states_path,
+            params_and_states_encoder_reference.forward_states);
+        read_vector_from_csv(
+            test_saving_paths_decoder.forward_states_path,
+            params_and_states_decoder_reference.forward_states);
+
+        // Compare the saved forward hidden states with the ones we got
+        if (!compare_vectors(params_and_states_encoder_reference.forward_states,
+                             params_and_states_encoder.forward_states, data,
+                             "encoder forward hidden states")) {
+            std::cout << "\033[1;31mTest for encoder FORWARD HIDDEN STATES has "
+                         "FAILED in " +
+                             data + " data\033[0m\n"
+                      << std::endl;
+            return false;
+        }
+
+        // Compare the saved forward hidden states with the ones we got
+        if (!compare_vectors(params_and_states_decoder_reference.forward_states,
+                             params_and_states_decoder.forward_states, data,
+                             "decoder forward hidden states")) {
+            std::cout << "\033[1;31mTest for decoder FORWARD HIDDEN STATES has "
+                         "FAILED in " +
+                             data + " data\033[0m\n"
+                      << std::endl;
+            return false;
+        }
+
+        // Read the saved backward hidden states reference
+        for (int i = 0; i < 2 * (net_prop_e.layers.size() - 2); i++)
+            params_and_states_encoder_reference.backward_states.push_back(
+                new std::vector<float>());
+        for (int i = 0; i < 2 * (net_prop_d.layers.size() - 2); i++)
+            params_and_states_decoder_reference.backward_states.push_back(
+                new std::vector<float>());
+
+        read_vector_from_csv(
+            test_saving_paths_encoder.backward_states_path,
+            params_and_states_encoder_reference.backward_states);
+        read_vector_from_csv(
+            test_saving_paths_decoder.backward_states_path,
+            params_and_states_decoder_reference.backward_states);
+
+        // Compare the saved backward hidden states with the ones we got
+        if (!compare_vectors(
+                params_and_states_encoder_reference.backward_states,
+                params_and_states_encoder.backward_states, data,
+                "encoder backward hidden states")) {
+            std::cout
+                << "\033[1;31mTest for encoder BACKWARD HIDDEN STATES has "
+                   "FAILED in " +
+                       data + " data\033[0m\n"
+                << std::endl;
+            return false;
+        }
+
+        // Compare the saved backward hidden states with the ones we got
+        if (!compare_vectors(
+                params_and_states_decoder_reference.backward_states,
+                params_and_states_decoder.backward_states, data,
+                "decoder backward hidden states")) {
+            std::cout
+                << "\033[1;31mTest for decoder BACKWARD HIDDEN STATES has "
+                   "FAILED in " +
+                       data + " data\033[0m\n"
+                << std::endl;
+            return false;
+        }
+    }
+    return true;
+}
diff --git a/test/autoencoder/test_autoencoder_gpu.cuh b/test/autoencoder/test_autoencoder_gpu.cuh
new file mode 100644
index 000000000..5b532b5ea
--- /dev/null
+++ b/test/autoencoder/test_autoencoder_gpu.cuh
@@ -0,0 +1,54 @@
+///////////////////////////////////////////////////////////////////////////////
+// File:         test_autoencoder_gpu.cuh
+// Description:  Header file for autoencoder test
+// Authors:      Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet
+// Created:      April 11, 2023
+// Updated:      April 12, 2023
+// Contact:      miquelflorensa11@gmail.com, luongha.nguyen@gmail.com &
+//               james.goulet@polymtl.ca
+// Copyright (c) 2023 Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet.
+// Some rights reserved.
+///////////////////////////////////////////////////////////////////////////////
+
+#pragma once
+
+#include <algorithm>
+#include <chrono>
+#include <cstdlib>
+#include <ctime>
+#include <iomanip>
+#include <iostream>
+#include <stdexcept>
+#include <string>
+
+#include "../../include/common.h"
+#include "../../include/cost.h"
+#include "../../include/dataloader.h"
+#include "../../include/derivative_calcul.cuh"
+#include "../../include/feed_forward.cuh"
+#include "../../include/global_param_update.cuh"
+#include "../../include/indices.h"
+#include "../../include/net_init.h"
+#include "../../include/net_prop.h"
+#include "../../include/param_feed_backward.cuh"
+#include "../../include/state_feed_backward.cuh"
+#include "../../include/struct_var.h"
+#include "../../include/tagi_network.cuh"
+#include "../../include/task.cuh"
+#include "../../include/utils.h"
+#include "../test_autoencoder.cuh"
+#include "../test_dataloader.h"
+#include "../test_utils.h"
+
+/**
+ * @brief Test the autoencoder network
+ *
+ * @param recompute_outputs indicates if the outputs should be recomputed
+ * @param date date of the test
+ * @param arch architecture of the network
+ * @param data dataset used
+ *
+ * @return true if the test was successful
+ */
+bool test_autoencoder_gpu(bool recompute_outputs, std::string date,
+                          std::string arch, std::string data);
diff --git a/test/cnn/test_cnn_gpu.cu b/test/cnn/test_cnn_gpu.cu
index 813ccd441..2680771f0 100644
--- a/test/cnn/test_cnn_gpu.cu
+++ b/test/cnn/test_cnn_gpu.cu
@@ -65,46 +65,14 @@ bool test_cnn_gpu(bool recompute_outputs, std::string date, std::string arch,
     path.curr_path = get_current_dir();
     std::string data_path = path.curr_path + "/test/data/" + data;
     std::string data_dir = path.curr_path + "/test/" + arch + "/data/";
-    std::string init_param_path_w =
-        data_dir + date + "_" + data + "_init_param_weights_w.csv";
-    std::string init_param_path_w_sc =
-        data_dir + date + "_" + data + "_init_param_weights_w_sc.csv";
-    std::string init_param_path_b =
-        data_dir + date + "_" + data + "_init_param_bias_b.csv";
-    std::string init_param_path_b_sc =
-        data_dir + date + "_" + data + "_init_param_bias_b_sc.csv";
-    std::string opt_param_path_w =
-        data_dir + date + "_" + data + "_opt_param_weights_w.csv";
-    std::string opt_param_path_w_sc =
-        data_dir + date + "_" + data + "_opt_param_weights_w_sc.csv";
-    std::string opt_param_path_b =
-        data_dir + date + "_" + data + "_opt_param_bias_b.csv";
-    std::string opt_param_path_b_sc =
-        data_dir + date + "_" + data + "_opt_param_bias_b_sc.csv";
-    std::string forward_states_path = data_dir + date +
-                                      "_forward_hidden_states_" + arch + "_" +
-                                      data + ".csv";
-    std::string backward_states_path = data_dir + date +
-                                       "_backward_hidden_states_" + arch + "_" +
-                                       data + ".csv";
+
+    TestSavingPaths test_saving_paths(path.curr_path, arch, data, date);
 
     // Train data
     ImageData imdb = image_dataloader(data, data_path, MU, SIGMA, NUM_CLASSES,
                                       tagi_net.prop);
 
-    std::vector<std::vector<float> *> weights;
-    weights.push_back(&tagi_net.theta.mw);
-    weights.push_back(&tagi_net.theta.Sw);
-    std::vector<std::vector<float> *> weights_sc;
-    weights_sc.push_back(&tagi_net.theta.mw_sc);
-    weights_sc.push_back(&tagi_net.theta.Sw_sc);
-
-    std::vector<std::vector<float> *> bias;
-    bias.push_back(&tagi_net.theta.mb);
-    bias.push_back(&tagi_net.theta.Sb);
-    std::vector<std::vector<float> *> bias_sc;
-    bias_sc.push_back(&tagi_net.theta.mb_sc);
-    bias_sc.push_back(&tagi_net.theta.Sb_sc);
+    TestParamAndStates params_and_states(tagi_net);
 
     // If we want to test but no data is available, we throw an error
     if (!recompute_outputs && !directory_exists(data_dir)) {
@@ -118,17 +86,11 @@ bool test_cnn_gpu(bool recompute_outputs, std::string date, std::string arch,
             return false;
         }
 
-        write_vector_to_csv(init_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(init_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-        write_vector_to_csv(init_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(init_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, true);
     }
 
     // Read the initial parameters (see test_utils.cpp for more details)
-    read_vector_from_csv(init_param_path_w, weights);
-    read_vector_from_csv(init_param_path_w_sc, weights_sc);
-    read_vector_from_csv(init_param_path_b, bias);
-    read_vector_from_csv(init_param_path_b_sc, bias_sc);
+    params_and_states.read_params(test_saving_paths, true);
 
     tagi_net.theta_gpu.copy_host_to_device();
 
@@ -138,75 +100,52 @@ bool test_cnn_gpu(bool recompute_outputs, std::string date, std::string arch,
     tagi_net.theta_gpu.copy_device_to_host();
     tagi_net.d_state_gpu.copy_device_to_host();
 
-    std::vector<std::vector<float> *> forward_states;
-    forward_states.push_back(&tagi_net.state.mz);
-    forward_states.push_back(&tagi_net.state.Sz);
-    forward_states.push_back(&tagi_net.state.ma);
-    forward_states.push_back(&tagi_net.state.Sa);
-    forward_states.push_back(&tagi_net.state.J);
+    add_forward_states(params_and_states.forward_states, tagi_net);
 
     std::vector<std::vector<float>> backward_states;
     std::string backward_states_header = "";
 
-    for (int i = 0; i < net.layers.size() - 2; i++) {
-        backward_states_header +=
-            "mean_" + std::to_string(i) + ",sigma_" + std::to_string(i) + ",";
-        backward_states.push_back(
-            std::get<0>(tagi_net.get_inovation_mean_var(i)));
-        backward_states.push_back(
-            std::get<1>(tagi_net.get_inovation_mean_var(i)));
-    }
+    add_backward_states(backward_states, backward_states_header, tagi_net,
+                        net.layers.size());
 
-    std::vector<std::vector<float> *> backward_states_ptr;
     for (int i = 0; i < backward_states.size(); i++)
-        backward_states_ptr.push_back(&backward_states[i]);
+        params_and_states.backward_states.push_back(&backward_states[i]);
 
     if (recompute_outputs) {
         // RESET OUPUTS
 
         // Write the parameters and hidden states
-        write_vector_to_csv(opt_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(opt_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-        write_vector_to_csv(opt_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(opt_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, false);
 
         // Write the forward hidden states
-        write_vector_to_csv(forward_states_path, "mz,Sz,ma,Sa,J",
-                            forward_states);
+        write_vector_to_csv(test_saving_paths.forward_states_path,
+                            "mz,Sz,ma,Sa,J", params_and_states.forward_states);
 
         // Write the backward hidden states
-        write_vector_to_csv(backward_states_path, backward_states_header,
-                            backward_states_ptr);
+        write_vector_to_csv(test_saving_paths.backward_states_path,
+                            backward_states_header,
+                            params_and_states.backward_states);
 
     } else {
         // PERFORM TESTS
 
         // Read the saved reference parameters
-        std::vector<std::vector<float> *> ref_weights;
-        std::vector<std::vector<float> *> ref_weights_sc;
-        std::vector<std::vector<float> *> ref_bias;
-        std::vector<std::vector<float> *> ref_bias_sc;
-
-        for (int i = 0; i < 2; i++) {
-            ref_weights.push_back(new std::vector<float>());
-            ref_weights_sc.push_back(new std::vector<float>());
-            ref_bias.push_back(new std::vector<float>());
-            ref_bias_sc.push_back(new std::vector<float>());
-        }
+        TestParamAndStates params_and_states_reference(tagi_net);
 
-        read_vector_from_csv(opt_param_path_w, ref_weights);
-        read_vector_from_csv(opt_param_path_w_sc, ref_weights_sc);
-        read_vector_from_csv(opt_param_path_b, ref_bias);
-        read_vector_from_csv(opt_param_path_b_sc, ref_bias_sc);
+        params_and_states_reference.read_params(test_saving_paths, false);
 
         tagi_net.theta_gpu.copy_host_to_device();
 
         // Compare optimal values with the ones we got
-        if (!compare_vectors(ref_weights, weights, data, "cnn weights") ||
-            !compare_vectors(ref_weights_sc, weights_sc, data,
+        if (!compare_vectors(params_and_states_reference.weights,
+                             params_and_states.weights, data, "cnn weights") ||
+            !compare_vectors(params_and_states_reference.weights_sc,
+                             params_and_states.weights_sc, data,
                              "cnn weights for residual network") ||
-            !compare_vectors(ref_bias, bias, data, "cnn bias") ||
-            !compare_vectors(ref_bias_sc, bias_sc, data,
+            !compare_vectors(params_and_states_reference.bias,
+                             params_and_states.bias, data, "cnn bias") ||
+            !compare_vectors(params_and_states_reference.bias_sc,
+                             params_and_states.bias_sc, data,
                              "cnn bias for residual network")) {
             std::cout << "\033[1;31mTest for CNN PARAMS has FAILED in " + data +
                              " data\033[0m\n"
@@ -215,14 +154,16 @@ bool test_cnn_gpu(bool recompute_outputs, std::string date, std::string arch,
         }
 
         // Read the saved forward hidden states reference
-        std::vector<std::vector<float> *> ref_forward_states;
         for (int i = 0; i < 5; i++)
-            ref_forward_states.push_back(new std::vector<float>());
+            params_and_states_reference.forward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(forward_states_path, ref_forward_states);
+        read_vector_from_csv(test_saving_paths.forward_states_path,
+                             params_and_states_reference.forward_states);
 
         // Compare the saved forward hidden states with the ones we got
-        if (!compare_vectors(ref_forward_states, forward_states, data,
+        if (!compare_vectors(params_and_states_reference.forward_states,
+                             params_and_states.forward_states, data,
                              "cnn forward hidden states")) {
             std::cout << "\033[1;31mTest for CNN FORWARD HIDDEN STATES has "
                          "FAILED in " +
@@ -232,14 +173,16 @@ bool test_cnn_gpu(bool recompute_outputs, std::string date, std::string arch,
         }
 
         // Read the saved backward hidden states reference
-        std::vector<std::vector<float> *> ref_backward_states;
         for (int i = 0; i < 2 * (net.layers.size() - 2); i++)
-            ref_backward_states.push_back(new std::vector<float>());
+            params_and_states_reference.backward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(backward_states_path, ref_backward_states);
+        read_vector_from_csv(test_saving_paths.backward_states_path,
+                             params_and_states_reference.backward_states);
 
         // Compare the saved backward hidden states with the ones we got
-        if (!compare_vectors(ref_backward_states, backward_states_ptr, data,
+        if (!compare_vectors(params_and_states_reference.backward_states,
+                             params_and_states.backward_states, data,
                              "cnn backward hidden states")) {
             std::cout << "\033[1;31mTest for CNN BACKWARD HIDDEN STATES has "
                          "FAILED in " +
@@ -249,4 +192,4 @@ bool test_cnn_gpu(bool recompute_outputs, std::string date, std::string arch,
         }
     }
     return true;
-}
\ No newline at end of file
+}
diff --git a/test/cnn_batch_norm/test_cnn_batch_norm_gpu.cu b/test/cnn_batch_norm/test_cnn_batch_norm_gpu.cu
index c7858ae31..62e5feb2c 100644
--- a/test/cnn_batch_norm/test_cnn_batch_norm_gpu.cu
+++ b/test/cnn_batch_norm/test_cnn_batch_norm_gpu.cu
@@ -3,7 +3,7 @@
 // Description:  Script to test the CNN batch norm GPU implementation of cuTAGI
 // Authors:      Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet
 // Created:      April 5, 2023
-// Updated:      April 6, 2023
+// Updated:      April 13, 2023
 // Contact:      miquelflorensa11@gmail.com & luongha.nguyen@gmail.com &
 //               james.goulet@polymtl.ca
 // Copyright (c) 2023 Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet.
@@ -65,46 +65,14 @@ bool test_cnn_batch_norm_gpu(bool recompute_outputs, std::string date,
     path.curr_path = get_current_dir();
     std::string data_path = path.curr_path + "/test/data/" + data;
     std::string data_dir = path.curr_path + "/test/" + arch + "/data/";
-    std::string init_param_path_w =
-        data_dir + date + "_" + data + "_init_param_weights_w.csv";
-    std::string init_param_path_w_sc =
-        data_dir + date + "_" + data + "_init_param_weights_w_sc.csv";
-    std::string init_param_path_b =
-        data_dir + date + "_" + data + "_init_param_bias_b.csv";
-    std::string init_param_path_b_sc =
-        data_dir + date + "_" + data + "_init_param_bias_b_sc.csv";
-    std::string opt_param_path_w =
-        data_dir + date + "_" + data + "_opt_param_weights_w.csv";
-    std::string opt_param_path_w_sc =
-        data_dir + date + "_" + data + "_opt_param_weights_w_sc.csv";
-    std::string opt_param_path_b =
-        data_dir + date + "_" + data + "_opt_param_bias_b.csv";
-    std::string opt_param_path_b_sc =
-        data_dir + date + "_" + data + "_opt_param_bias_b_sc.csv";
-    std::string forward_states_path = data_dir + date +
-                                      "_forward_hidden_states_" + arch + "_" +
-                                      data + ".csv";
-    std::string backward_states_path = data_dir + date +
-                                       "_backward_hidden_states_" + arch + "_" +
-                                       data + ".csv";
+
+    TestSavingPaths test_saving_paths(path.curr_path, arch, data, date);
 
     // Train data
     ImageData imdb = image_dataloader(data, data_path, MU, SIGMA, NUM_CLASSES,
                                       tagi_net.prop);
 
-    std::vector<std::vector<float> *> weights;
-    weights.push_back(&tagi_net.theta.mw);
-    weights.push_back(&tagi_net.theta.Sw);
-    std::vector<std::vector<float> *> weights_sc;
-    weights_sc.push_back(&tagi_net.theta.mw_sc);
-    weights_sc.push_back(&tagi_net.theta.Sw_sc);
-
-    std::vector<std::vector<float> *> bias;
-    bias.push_back(&tagi_net.theta.mb);
-    bias.push_back(&tagi_net.theta.Sb);
-    std::vector<std::vector<float> *> bias_sc;
-    bias_sc.push_back(&tagi_net.theta.mb_sc);
-    bias_sc.push_back(&tagi_net.theta.Sb_sc);
+    TestParamAndStates params_and_states(tagi_net);
 
     // If we want to test but no data is available, we throw an error
     if (!recompute_outputs && !directory_exists(data_dir)) {
@@ -118,17 +86,11 @@ bool test_cnn_batch_norm_gpu(bool recompute_outputs, std::string date,
             return false;
         }
 
-        write_vector_to_csv(init_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(init_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-        write_vector_to_csv(init_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(init_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, true);
     }
 
     // Read the initial parameters (see test_utils.cpp for more details)
-    read_vector_from_csv(init_param_path_w, weights);
-    read_vector_from_csv(init_param_path_w_sc, weights_sc);
-    read_vector_from_csv(init_param_path_b, bias);
-    read_vector_from_csv(init_param_path_b_sc, bias_sc);
+    params_and_states.read_params(test_saving_paths, true);
 
     tagi_net.theta_gpu.copy_host_to_device();
 
@@ -138,96 +100,76 @@ bool test_cnn_batch_norm_gpu(bool recompute_outputs, std::string date,
     tagi_net.theta_gpu.copy_device_to_host();
     tagi_net.d_state_gpu.copy_device_to_host();
 
-    std::vector<std::vector<float> *> forward_states;
-    forward_states.push_back(&tagi_net.state.mz);
-    forward_states.push_back(&tagi_net.state.Sz);
-    forward_states.push_back(&tagi_net.state.ma);
-    forward_states.push_back(&tagi_net.state.Sa);
-    forward_states.push_back(&tagi_net.state.J);
+    add_forward_states(params_and_states.forward_states, tagi_net);
 
     std::vector<std::vector<float>> backward_states;
     std::string backward_states_header = "";
 
-    for (int i = 0; i < net.layers.size() - 2; i++) {
-        backward_states_header +=
-            "mean_" + std::to_string(i) + ",sigma_" + std::to_string(i) + ",";
-        backward_states.push_back(
-            std::get<0>(tagi_net.get_inovation_mean_var(i)));
-        backward_states.push_back(
-            std::get<1>(tagi_net.get_inovation_mean_var(i)));
-    }
+    add_backward_states(backward_states, backward_states_header, tagi_net,
+                        net.layers.size());
 
-    std::vector<std::vector<float> *> backward_states_ptr;
     for (int i = 0; i < backward_states.size(); i++)
-        backward_states_ptr.push_back(&backward_states[i]);
+        params_and_states.backward_states.push_back(&backward_states[i]);
 
     if (recompute_outputs) {
         // RESET OUPUTS
 
         // Write the parameters and hidden states
-        write_vector_to_csv(opt_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(opt_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-        write_vector_to_csv(opt_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(opt_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, false);
 
         // Write the forward hidden states
-        write_vector_to_csv(forward_states_path, "mz,Sz,ma,Sa,J",
-                            forward_states);
+        write_vector_to_csv(test_saving_paths.forward_states_path,
+                            "mz,Sz,ma,Sa,J", params_and_states.forward_states);
 
         // Write the backward hidden states
-        write_vector_to_csv(backward_states_path, backward_states_header,
-                            backward_states_ptr);
+        write_vector_to_csv(test_saving_paths.backward_states_path,
+                            backward_states_header,
+                            params_and_states.backward_states);
 
     } else {
         // PERFORM TESTS
 
         // Read the saved reference parameters
-        std::vector<std::vector<float> *> ref_weights;
-        std::vector<std::vector<float> *> ref_weights_sc;
-        std::vector<std::vector<float> *> ref_bias;
-        std::vector<std::vector<float> *> ref_bias_sc;
-
-        for (int i = 0; i < 2; i++) {
-            ref_weights.push_back(new std::vector<float>());
-            ref_weights_sc.push_back(new std::vector<float>());
-            ref_bias.push_back(new std::vector<float>());
-            ref_bias_sc.push_back(new std::vector<float>());
-        }
+        TestParamAndStates params_and_states_reference(tagi_net);
 
-        read_vector_from_csv(opt_param_path_w, ref_weights);
-        read_vector_from_csv(opt_param_path_w_sc, ref_weights_sc);
-        read_vector_from_csv(opt_param_path_b, ref_bias);
-        read_vector_from_csv(opt_param_path_b_sc, ref_bias_sc);
+        params_and_states_reference.read_params(test_saving_paths, false);
 
         tagi_net.theta_gpu.copy_host_to_device();
 
         // Compare optimal values with the ones we got
-        if (!compare_vectors(ref_weights, weights, data,
+        if (!compare_vectors(params_and_states_reference.weights,
+                             params_and_states.weights, data,
                              "cnn batch norm. weights") ||
-            !compare_vectors(ref_weights_sc, weights_sc, data,
+            !compare_vectors(params_and_states_reference.weights_sc,
+                             params_and_states.weights_sc, data,
                              "cnn batch norm. weights for residual network") ||
-            !compare_vectors(ref_bias, bias, data, "cnn batch norm. bias") ||
-            !compare_vectors(ref_bias_sc, bias_sc, data,
+            !compare_vectors(params_and_states_reference.bias,
+                             params_and_states.bias, data,
+                             "cnn batch norm. bias") ||
+            !compare_vectors(params_and_states_reference.bias_sc,
+                             params_and_states.bias_sc, data,
                              "cnn batch norm. bias for residual network")) {
-            std::cout << "\033[1;31mTest for CNN BATCH NORMALIZATION PARAMS "
-                         "has FAILED in " +
-                             data + " data\033[0m\n"
-                      << std::endl;
+            std::cout
+                << "\033[1;31mTest for CNN batch norm. PARAMS has FAILED in " +
+                       data + " data\033[0m\n"
+                << std::endl;
             return false;
         }
 
         // Read the saved forward hidden states reference
-        std::vector<std::vector<float> *> ref_forward_states;
         for (int i = 0; i < 5; i++)
-            ref_forward_states.push_back(new std::vector<float>());
+            params_and_states_reference.forward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(forward_states_path, ref_forward_states);
+        read_vector_from_csv(test_saving_paths.forward_states_path,
+                             params_and_states_reference.forward_states);
 
         // Compare the saved forward hidden states with the ones we got
-        if (!compare_vectors(ref_forward_states, forward_states, data,
+        if (!compare_vectors(params_and_states_reference.forward_states,
+                             params_and_states.forward_states, data,
                              "cnn batch norm. forward hidden states")) {
-            std::cout << "\033[1;31mTest for CNN BATCH NORMALIZATION FORWARD "
-                         "HIDDEN STATES has "
+            std::cout << "\033[1;31mTest for CNN batch norm. FORWARD HIDDEN "
+                         "STATES has "
                          "FAILED in " +
                              data + " data\033[0m\n"
                       << std::endl;
@@ -235,17 +177,19 @@ bool test_cnn_batch_norm_gpu(bool recompute_outputs, std::string date,
         }
 
         // Read the saved backward hidden states reference
-        std::vector<std::vector<float> *> ref_backward_states;
         for (int i = 0; i < 2 * (net.layers.size() - 2); i++)
-            ref_backward_states.push_back(new std::vector<float>());
+            params_and_states_reference.backward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(backward_states_path, ref_backward_states);
+        read_vector_from_csv(test_saving_paths.backward_states_path,
+                             params_and_states_reference.backward_states);
 
         // Compare the saved backward hidden states with the ones we got
-        if (!compare_vectors(ref_backward_states, backward_states_ptr, data,
+        if (!compare_vectors(params_and_states_reference.backward_states,
+                             params_and_states.backward_states, data,
                              "cnn batch norm. backward hidden states")) {
-            std::cout << "\033[1;31mTest for CNN BATCH NORMALIZATION BACKWARD "
-                         "HIDDEN STATES has "
+            std::cout << "\033[1;31mTest for CNN batch norm. BACKWARD HIDDEN "
+                         "STATES has "
                          "FAILED in " +
                              data + " data\033[0m\n"
                       << std::endl;
@@ -253,4 +197,4 @@ bool test_cnn_batch_norm_gpu(bool recompute_outputs, std::string date,
         }
     }
     return true;
-}
\ No newline at end of file
+}
diff --git a/test/data/last_dates.csv b/test/data/last_dates.csv
index aa071ce03..e9aac97ed 100644
--- a/test/data/last_dates.csv
+++ b/test/data/last_dates.csv
@@ -1,2 +1,2 @@
 fnn,fnn_heteros,fnn_full_cov,fnn_derivatives,lstm,act_func,cnn,cnn_batch_norm,autoencoder
-2023_04_06,2023_04_06,2023_04_06,2023_04_06,2023_04_06,2023_04_06,2023_04_06,2023_04_06,2023_03_28
+2023_04_17,2023_04_17,2023_04_17,2023_04_17,2023_04_17,2023_04_17,2023_04_17,2023_04_17,2023_04_17
diff --git a/test/fnn/test_fnn_cpu.cpp b/test/fnn/test_fnn_cpu.cpp
index f956901d7..e27479bf3 100644
--- a/test/fnn/test_fnn_cpu.cpp
+++ b/test/fnn/test_fnn_cpu.cpp
@@ -47,28 +47,8 @@ bool test_fnn_cpu(bool recompute_outputs, std::string date, std::string arch,
     path.curr_path = get_current_dir();
     std::string data_path = path.curr_path + "/test/data/" + data;
     std::string data_dir = path.curr_path + "/test/" + arch + "/data/";
-    std::string init_param_path_w =
-        data_dir + date + "_" + data + "_init_param_weights_w.csv";
-    std::string init_param_path_w_sc =
-        data_dir + date + "_" + data + "_init_param_weights_w_sc.csv";
-    std::string init_param_path_b =
-        data_dir + date + "_" + data + "_init_param_bias_b.csv";
-    std::string init_param_path_b_sc =
-        data_dir + date + "_" + data + "_init_param_bias_b_sc.csv";
-    std::string opt_param_path_w =
-        data_dir + date + "_" + data + "_opt_param_weights_w.csv";
-    std::string opt_param_path_w_sc =
-        data_dir + date + "_" + data + "_opt_param_weights_w_sc.csv";
-    std::string opt_param_path_b =
-        data_dir + date + "_" + data + "_opt_param_bias_b.csv";
-    std::string opt_param_path_b_sc =
-        data_dir + date + "_" + data + "_opt_param_bias_b_sc.csv";
-    std::string forward_states_path = data_dir + date +
-                                      "_forward_hidden_states_" + arch + "_" +
-                                      data + ".csv";
-    std::string backward_states_path = data_dir + date +
-                                       "_backward_hidden_states_" + arch + "_" +
-                                       data + ".csv";
+
+    TestSavingPaths test_saving_paths(path.curr_path, arch, data, date);
 
     // Train data
     Dataloader train_db = train_data(data, tagi_net, data_path, NORMALIZE);
@@ -76,19 +56,7 @@ bool test_fnn_cpu(bool recompute_outputs, std::string date, std::string arch,
     Dataloader test_db =
         test_data(data, tagi_net, data_path, train_db, NORMALIZE);
 
-    std::vector<std::vector<float> *> weights;
-    weights.push_back(&tagi_net.theta.mw);
-    weights.push_back(&tagi_net.theta.Sw);
-    std::vector<std::vector<float> *> weights_sc;
-    weights_sc.push_back(&tagi_net.theta.mw_sc);
-    weights_sc.push_back(&tagi_net.theta.Sw_sc);
-
-    std::vector<std::vector<float> *> bias;
-    bias.push_back(&tagi_net.theta.mb);
-    bias.push_back(&tagi_net.theta.Sb);
-    std::vector<std::vector<float> *> bias_sc;
-    bias_sc.push_back(&tagi_net.theta.mb_sc);
-    bias_sc.push_back(&tagi_net.theta.Sb_sc);
+    TestParamAndStates params_and_states(tagi_net);
 
     // If we want to test but no data is available, we throw an error
     if (!recompute_outputs && !directory_exists(data_dir)) {
@@ -101,89 +69,59 @@ bool test_fnn_cpu(bool recompute_outputs, std::string date, std::string arch,
             std::cout << "Error: could not create data directory" << std::endl;
             return false;
         }
-        write_vector_to_csv(init_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(init_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-
-        write_vector_to_csv(init_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(init_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, true);
     }
 
     // Read the initial parameters (see tes_utils.cpp for more details)
-    read_vector_from_csv(init_param_path_w, weights);
-    read_vector_from_csv(init_param_path_w_sc, weights_sc);
-    read_vector_from_csv(init_param_path_b, bias);
-    read_vector_from_csv(init_param_path_b_sc, bias_sc);
+    params_and_states.read_params(test_saving_paths, true);
 
     // Train the network
     regression_train(tagi_net, train_db, EPOCHS);
 
-    std::vector<std::vector<float> *> forward_states;
-    forward_states.push_back(&tagi_net.state.mz);
-    forward_states.push_back(&tagi_net.state.Sz);
-    forward_states.push_back(&tagi_net.state.ma);
-    forward_states.push_back(&tagi_net.state.Sa);
-    forward_states.push_back(&tagi_net.state.J);
+    add_forward_states(params_and_states.forward_states, tagi_net);
 
     std::vector<std::vector<float>> backward_states;
     std::string backward_states_header = "";
 
-    for (int i = 0; i < net.layers.size() - 2; i++) {
-        backward_states_header +=
-            "mean_" + std::to_string(i) + ",sigma_" + std::to_string(i) + ",";
-        backward_states.push_back(
-            std::get<0>(tagi_net.get_inovation_mean_var(i)));
-        backward_states.push_back(
-            std::get<1>(tagi_net.get_inovation_mean_var(i)));
-    }
+    add_backward_states(backward_states, backward_states_header, tagi_net,
+                        net.layers.size());
 
-    std::vector<std::vector<float> *> backward_states_ptr;
     for (int i = 0; i < backward_states.size(); i++)
-        backward_states_ptr.push_back(&backward_states[i]);
+        params_and_states.backward_states.push_back(&backward_states[i]);
 
     if (recompute_outputs) {
         // RESET OUPUTS
 
         // Write the parameters and hidden states
-        write_vector_to_csv(opt_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(opt_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-        write_vector_to_csv(opt_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(opt_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, false);
 
         // Write the forward hidden states
-        write_vector_to_csv(forward_states_path, "mz,Sz,ma,Sa,J",
-                            forward_states);
+        write_vector_to_csv(test_saving_paths.forward_states_path,
+                            "mz,Sz,ma,Sa,J", params_and_states.forward_states);
 
         // Write the backward hidden states
-        write_vector_to_csv(backward_states_path, backward_states_header,
-                            backward_states_ptr);
+        write_vector_to_csv(test_saving_paths.backward_states_path,
+                            backward_states_header,
+                            params_and_states.backward_states);
 
     } else {
         // PERFORM TESTS
 
         // Read the saved reference parameters
-        std::vector<std::vector<float> *> ref_weights;
-        std::vector<std::vector<float> *> ref_weights_sc;
-        std::vector<std::vector<float> *> ref_bias;
-        std::vector<std::vector<float> *> ref_bias_sc;
-
-        for (int i = 0; i < 2; i++) {
-            ref_weights.push_back(new std::vector<float>());
-            ref_weights_sc.push_back(new std::vector<float>());
-            ref_bias.push_back(new std::vector<float>());
-            ref_bias_sc.push_back(new std::vector<float>());
-        }
+        TestParamAndStates params_and_states_reference(tagi_net);
 
-        read_vector_from_csv(opt_param_path_w, ref_weights);
-        read_vector_from_csv(opt_param_path_w_sc, ref_weights_sc);
-        read_vector_from_csv(opt_param_path_b, ref_bias);
-        read_vector_from_csv(opt_param_path_b_sc, ref_bias_sc);
+        params_and_states_reference.read_params(test_saving_paths, false);
 
         // Compare optimal values with the ones we got
-        if (!compare_vectors(ref_weights, weights, data, "fnn weights") ||
-            !compare_vectors(ref_weights_sc, weights_sc, data,
+        if (!compare_vectors(params_and_states_reference.weights,
+                             params_and_states.weights, data, "fnn weights") ||
+            !compare_vectors(params_and_states_reference.weights_sc,
+                             params_and_states.weights_sc, data,
                              "fnn weights for residual network") ||
-            !compare_vectors(ref_bias, bias, data, "fnn bias") ||
-            !compare_vectors(ref_bias_sc, bias_sc, data,
+            !compare_vectors(params_and_states_reference.bias,
+                             params_and_states.bias, data, "fnn bias") ||
+            !compare_vectors(params_and_states_reference.bias_sc,
+                             params_and_states.bias_sc, data,
                              "fnn bias for residual network")) {
             std::cout << "\033[1;31mTest for FNN PARAMS has FAILED in " + data +
                              " data\033[0m\n"
@@ -192,14 +130,16 @@ bool test_fnn_cpu(bool recompute_outputs, std::string date, std::string arch,
         }
 
         // Read the saved forward hidden states reference
-        std::vector<std::vector<float> *> ref_forward_states;
         for (int i = 0; i < 5; i++)
-            ref_forward_states.push_back(new std::vector<float>());
+            params_and_states_reference.forward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(forward_states_path, ref_forward_states);
+        read_vector_from_csv(test_saving_paths.forward_states_path,
+                             params_and_states_reference.forward_states);
 
         // Compare the saved forward hidden states with the ones we got
-        if (!compare_vectors(ref_forward_states, forward_states, data,
+        if (!compare_vectors(params_and_states_reference.forward_states,
+                             params_and_states.forward_states, data,
                              "fnn forward hidden states")) {
             std::cout << "\033[1;31mTest for FNN FORWARD HIDDEN STATES has "
                          "FAILED in " +
@@ -209,14 +149,16 @@ bool test_fnn_cpu(bool recompute_outputs, std::string date, std::string arch,
         }
 
         // Read the saved backward hidden states reference
-        std::vector<std::vector<float> *> ref_backward_states;
         for (int i = 0; i < 2 * (net.layers.size() - 2); i++)
-            ref_backward_states.push_back(new std::vector<float>());
+            params_and_states_reference.backward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(backward_states_path, ref_backward_states);
+        read_vector_from_csv(test_saving_paths.backward_states_path,
+                             params_and_states_reference.backward_states);
 
         // Compare the saved backward hidden states with the ones we got
-        if (!compare_vectors(ref_backward_states, backward_states_ptr, data,
+        if (!compare_vectors(params_and_states_reference.backward_states,
+                             params_and_states.backward_states, data,
                              "fnn backward hidden states")) {
             std::cout << "\033[1;31mTest for FNN BACKWARD HIDDEN STATES has "
                          "FAILED in " +
@@ -226,4 +168,4 @@ bool test_fnn_cpu(bool recompute_outputs, std::string date, std::string arch,
         }
     }
     return true;
-}
\ No newline at end of file
+}
diff --git a/test/fnn_derivatives/test_fnn_derivatives_cpu.cpp b/test/fnn_derivatives/test_fnn_derivatives_cpu.cpp
index 52891d3e5..d7e3e4765 100644
--- a/test/fnn_derivatives/test_fnn_derivatives_cpu.cpp
+++ b/test/fnn_derivatives/test_fnn_derivatives_cpu.cpp
@@ -50,30 +50,8 @@ bool test_fnn_derivatives_cpu(bool recompute_outputs, std::string date,
     path.curr_path = get_current_dir();
     std::string data_path = path.curr_path + "/test/data/" + data;
     std::string data_dir = path.curr_path + "/test/" + arch + "/data/";
-    std::string init_param_path_w =
-        data_dir + date + "_" + data + "_init_param_weights_w.csv";
-    std::string init_param_path_w_sc =
-        data_dir + date + "_" + data + "_init_param_weights_w_sc.csv";
-    std::string init_param_path_b =
-        data_dir + date + "_" + data + "_init_param_bias_b.csv";
-    std::string init_param_path_b_sc =
-        data_dir + date + "_" + data + "_init_param_bias_b_sc.csv";
-    std::string opt_param_path_w =
-        data_dir + date + "_" + data + "_opt_param_weights_w.csv";
-    std::string opt_param_path_w_sc =
-        data_dir + date + "_" + data + "_opt_param_weights_w_sc.csv";
-    std::string opt_param_path_b =
-        data_dir + date + "_" + data + "_opt_param_bias_b.csv";
-    std::string opt_param_path_b_sc =
-        data_dir + date + "_" + data + "_opt_param_bias_b_sc.csv";
-    std::string forward_states_path = data_dir + date +
-                                      "_forward_hidden_states_" + arch + "_" +
-                                      data + ".csv";
-    std::string backward_states_path = data_dir + date +
-                                       "_backward_hidden_states_" + arch + "_" +
-                                       data + ".csv";
-    std::string input_derivative_path =
-        data_dir + date + "_input_derivative_" + arch + "_" + data + ".csv";
+
+    TestSavingPaths test_saving_paths(path.curr_path, arch, data, date);
 
     // Train data
     Dataloader train_db = train_data(data, tagi_net, data_path, NORMALIZE);
@@ -81,19 +59,7 @@ bool test_fnn_derivatives_cpu(bool recompute_outputs, std::string date,
     Dataloader test_db =
         test_data(data, tagi_net, data_path, train_db, NORMALIZE);
 
-    std::vector<std::vector<float> *> weights;
-    weights.push_back(&tagi_net.theta.mw);
-    weights.push_back(&tagi_net.theta.Sw);
-    std::vector<std::vector<float> *> weights_sc;
-    weights_sc.push_back(&tagi_net.theta.mw_sc);
-    weights_sc.push_back(&tagi_net.theta.Sw_sc);
-
-    std::vector<std::vector<float> *> bias;
-    bias.push_back(&tagi_net.theta.mb);
-    bias.push_back(&tagi_net.theta.Sb);
-    std::vector<std::vector<float> *> bias_sc;
-    bias_sc.push_back(&tagi_net.theta.mb_sc);
-    bias_sc.push_back(&tagi_net.theta.Sb_sc);
+    TestParamAndStates params_and_states(tagi_net);
 
     // If we want to test but no data is available, we throw an error
     if (!recompute_outputs && !directory_exists(data_dir)) {
@@ -106,157 +72,135 @@ bool test_fnn_derivatives_cpu(bool recompute_outputs, std::string date,
             std::cout << "Error: could not create data directory" << std::endl;
             return false;
         }
-        write_vector_to_csv(init_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(init_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-
-        write_vector_to_csv(init_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(init_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, true);
     }
 
     // Read the initial parameters (see tes_utils.cpp for more details)
-    read_vector_from_csv(init_param_path_w, weights);
-    read_vector_from_csv(init_param_path_w_sc, weights_sc);
-    read_vector_from_csv(init_param_path_b, bias);
-    read_vector_from_csv(init_param_path_b_sc, bias_sc);
+    params_and_states.read_params(test_saving_paths, true);
 
     // Train the network
     regression_train(tagi_net, train_db, EPOCHS);
 
     // Take pointers from all forward states
-    std::vector<std::vector<float> *> forward_states;
-    forward_states.push_back(&tagi_net.state.mz);
-    forward_states.push_back(&tagi_net.state.Sz);
-    forward_states.push_back(&tagi_net.state.ma);
-    forward_states.push_back(&tagi_net.state.Sa);
-    forward_states.push_back(&tagi_net.state.J);
-
-    // Take pointers from all backward states
+    add_forward_states(params_and_states.forward_states, tagi_net);
+
     std::vector<std::vector<float>> backward_states;
     std::string backward_states_header = "";
 
-    for (int i = 0; i < net.layers.size() - 2; i++) {
-        backward_states_header +=
-            "mean_" + std::to_string(i) + ",sigma_" + std::to_string(i) + ",";
-        backward_states.push_back(
-            std::get<0>(tagi_net.get_inovation_mean_var(i)));
-        backward_states.push_back(
-            std::get<1>(tagi_net.get_inovation_mean_var(i)));
-    }
+    add_backward_states(backward_states, backward_states_header, tagi_net,
+                        net.layers.size());
 
-    std::vector<std::vector<float> *> backward_states_ptr;
     for (int i = 0; i < backward_states.size(); i++)
-        backward_states_ptr.push_back(&backward_states[i]);
+        params_and_states.backward_states.push_back(&backward_states[i]);
 
     // Take pointers from all input derivatives
     std::vector<std::vector<float>> input_derivatives;
     input_derivatives.push_back(std::get<0>(tagi_net.get_derivatives(0)));
     input_derivatives.push_back(std::get<1>(tagi_net.get_derivatives(0)));
 
-    std::vector<std::vector<float> *> input_derivatives_ptr;
     for (int i = 0; i < input_derivatives.size(); i++)
-        input_derivatives_ptr.push_back(&input_derivatives[i]);
+        params_and_states.input_derivatives.push_back(&input_derivatives[i]);
 
     if (recompute_outputs) {
         // RESET OUPUTS
 
         // Write the parameters and hidden states
-        write_vector_to_csv(opt_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(opt_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-        write_vector_to_csv(opt_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(opt_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, false);
 
         // Write the forward hidden states
-        write_vector_to_csv(forward_states_path, "mz,Sz,ma,Sa,J",
-                            forward_states);
+        write_vector_to_csv(test_saving_paths.forward_states_path,
+                            "mz,Sz,ma,Sa,J", params_and_states.forward_states);
 
         // Write the backward hidden states
-        write_vector_to_csv(backward_states_path, backward_states_header,
-                            backward_states_ptr);
+        write_vector_to_csv(test_saving_paths.backward_states_path,
+                            backward_states_header,
+                            params_and_states.backward_states);
 
         // Write estimated derivatives of the input layer
-        write_vector_to_csv(input_derivative_path, "md,Sd",
-                            input_derivatives_ptr);
+        write_vector_to_csv(test_saving_paths.input_derivative_path, "md,Sd",
+                            params_and_states.input_derivatives);
 
     } else {
         // PERFORM TESTS
 
         // Read the saved reference parameters
-        std::vector<std::vector<float> *> ref_weights;
-        std::vector<std::vector<float> *> ref_weights_sc;
-        std::vector<std::vector<float> *> ref_bias;
-        std::vector<std::vector<float> *> ref_bias_sc;
-
-        for (int i = 0; i < 2; i++) {
-            ref_weights.push_back(new std::vector<float>());
-            ref_weights_sc.push_back(new std::vector<float>());
-            ref_bias.push_back(new std::vector<float>());
-            ref_bias_sc.push_back(new std::vector<float>());
-        }
+        TestParamAndStates params_and_states_reference(tagi_net);
 
-        read_vector_from_csv(opt_param_path_w, ref_weights);
-        read_vector_from_csv(opt_param_path_w_sc, ref_weights_sc);
-        read_vector_from_csv(opt_param_path_b, ref_bias);
-        read_vector_from_csv(opt_param_path_b_sc, ref_bias_sc);
+        params_and_states_reference.read_params(test_saving_paths, false);
 
         // Compare optimal values with the ones we got
-        if (!compare_vectors(ref_weights, weights, data,
-                             "fnn derivatives weights") ||
-            !compare_vectors(ref_weights_sc, weights_sc, data,
-                             "fnn derivatives for residual network") ||
-            !compare_vectors(ref_bias, bias, data, "fnn derivatives bias") ||
-            !compare_vectors(ref_bias_sc, bias_sc, data,
-                             "fnn derivatives bias for residual network")) {
-            std::cout
-                << "\033[1;31mTest for FNN DERIVATIVES PARAMS has FAILED in " +
-                       data + " data\033[0m\n"
-                << std::endl;
+        if (!compare_vectors(params_and_states_reference.weights,
+                             params_and_states.weights, data,
+                             "fnn input derivatives weights") ||
+            !compare_vectors(
+                params_and_states_reference.weights_sc,
+                params_and_states.weights_sc, data,
+                "fnn input derivatives weights for residual network") ||
+            !compare_vectors(params_and_states_reference.bias,
+                             params_and_states.bias, data,
+                             "fnn input derivatives bias") ||
+            !compare_vectors(
+                params_and_states_reference.bias_sc, params_and_states.bias_sc,
+                data, "fnn input derivatives bias for residual network")) {
+            std::cout << "\033[1;31mTest for FNN INPUT DERIVATIVES PARAMS has "
+                         "FAILED in " +
+                             data + " data\033[0m\n"
+                      << std::endl;
             return false;
         }
 
         // Read the saved forward hidden states reference
-        std::vector<std::vector<float> *> ref_forward_states;
         for (int i = 0; i < 5; i++)
-            ref_forward_states.push_back(new std::vector<float>());
+            params_and_states_reference.forward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(forward_states_path, ref_forward_states);
+        read_vector_from_csv(test_saving_paths.forward_states_path,
+                             params_and_states_reference.forward_states);
 
         // Compare the saved forward hidden states with the ones we got
-        if (!compare_vectors(ref_forward_states, forward_states, data,
-                             "fnn derivatives forward hidden states")) {
-            std::cout << "\033[1;31mTest for FNN DERIVATIVES FORWARD HIDDEN "
-                         "STATES has FAILED in " +
+        if (!compare_vectors(params_and_states_reference.forward_states,
+                             params_and_states.forward_states, data,
+                             "fnn input derivatives forward hidden states")) {
+            std::cout << "\033[1;31mTest for FNN INPUT DERIVATIVES FORWARD "
+                         "HIDDEN STATES has "
+                         "FAILED in " +
                              data + " data\033[0m\n"
                       << std::endl;
             return false;
         }
 
         // Read the saved backward hidden states reference
-        std::vector<std::vector<float> *> ref_backward_states;
         for (int i = 0; i < 2 * (net.layers.size() - 2); i++)
-            ref_backward_states.push_back(new std::vector<float>());
+            params_and_states_reference.backward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(backward_states_path, ref_backward_states);
+        read_vector_from_csv(test_saving_paths.backward_states_path,
+                             params_and_states_reference.backward_states);
 
         // Compare the saved backward hidden states with the ones we got
-        if (!compare_vectors(ref_backward_states, backward_states_ptr, data,
-                             "fnn derivatives backward hidden states")) {
-            std::cout << "\033[1;31mTest for FNN DERIVATIVES BACKWARD HIDDEN "
-                         "STATES has FAILED in " +
+        if (!compare_vectors(params_and_states_reference.backward_states,
+                             params_and_states.backward_states, data,
+                             "fnn input derivatives backward hidden states")) {
+            std::cout << "\033[1;31mTest for FNN INPUT DERIVATIVES BACKWARD "
+                         "HIDDEN STATES has "
+                         "FAILED in " +
                              data + " data\033[0m\n"
                       << std::endl;
             return false;
         }
 
         // Read the saved input derivative reference
-        std::vector<std::vector<float> *> ref_input_derivative;
         for (int i = 0; i < 2; i++)
-            ref_input_derivative.push_back(new std::vector<float>());
+            params_and_states_reference.input_derivatives.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(input_derivative_path, ref_input_derivative);
+        read_vector_from_csv(test_saving_paths.input_derivative_path,
+                             params_and_states_reference.input_derivatives);
 
         // Compare the saved input derivatives with the ones we got
-        if (!compare_vectors(ref_input_derivative, input_derivatives_ptr, data,
-                             "fnn derivatives input derivatives")) {
+        if (!compare_vectors(params_and_states_reference.input_derivatives,
+                             params_and_states.input_derivatives, data,
+                             "fnn input derivatives")) {
             std::cout << "\033[1;31mTest for FNN INPUT DERIVATIVES has "
                          "FAILED in " +
                              data + " data\033[0m\n"
@@ -265,4 +209,4 @@ bool test_fnn_derivatives_cpu(bool recompute_outputs, std::string date,
         }
     }
     return true;
-}
\ No newline at end of file
+}
diff --git a/test/fnn_full_cov/test_fnn_full_cov_cpu.cpp b/test/fnn_full_cov/test_fnn_full_cov_cpu.cpp
index 4e3a901a3..186352881 100644
--- a/test/fnn_full_cov/test_fnn_full_cov_cpu.cpp
+++ b/test/fnn_full_cov/test_fnn_full_cov_cpu.cpp
@@ -50,28 +50,8 @@ bool test_fnn_full_cov_cpu(bool recompute_outputs, std::string date,
     path.curr_path = get_current_dir();
     std::string data_path = path.curr_path + "/test/data/" + data;
     std::string data_dir = path.curr_path + "/test/" + arch + "/data/";
-    std::string init_param_path_w =
-        data_dir + date + "_" + data + "_init_param_weights_w.csv";
-    std::string init_param_path_w_sc =
-        data_dir + date + "_" + data + "_init_param_weights_w_sc.csv";
-    std::string init_param_path_b =
-        data_dir + date + "_" + data + "_init_param_bias_b.csv";
-    std::string init_param_path_b_sc =
-        data_dir + date + "_" + data + "_init_param_bias_b_sc.csv";
-    std::string opt_param_path_w =
-        data_dir + date + "_" + data + "_opt_param_weights_w.csv";
-    std::string opt_param_path_w_sc =
-        data_dir + date + "_" + data + "_opt_param_weights_w_sc.csv";
-    std::string opt_param_path_b =
-        data_dir + date + "_" + data + "_opt_param_bias_b.csv";
-    std::string opt_param_path_b_sc =
-        data_dir + date + "_" + data + "_opt_param_bias_b_sc.csv";
-    std::string forward_states_path = data_dir + date +
-                                      "_forward_hidden_states_" + arch + "_" +
-                                      data + ".csv";
-    std::string backward_states_path = data_dir + date +
-                                       "_backward_hidden_states_" + arch + "_" +
-                                       data + ".csv";
+
+    TestSavingPaths test_saving_paths(path.curr_path, arch, data, date);
 
     // Train data
     Dataloader train_db = train_data(data, tagi_net, data_path, NORMALIZE);
@@ -104,18 +84,21 @@ bool test_fnn_full_cov_cpu(bool recompute_outputs, std::string date,
             std::cout << "Error: could not create data directory" << std::endl;
             return false;
         }
-        write_vector_to_csv(init_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(init_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-
-        write_vector_to_csv(init_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(init_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        write_vector_to_csv(test_saving_paths.init_param_path_w, "mw,Sw",
+                            weights);
+        write_vector_to_csv(test_saving_paths.init_param_path_w_sc,
+                            "mw_sc,Sw_sc", weights_sc);
+
+        write_vector_to_csv(test_saving_paths.init_param_path_b, "mb,Sb", bias);
+        write_vector_to_csv(test_saving_paths.init_param_path_b_sc,
+                            "mb_sc,Sb_sc", bias_sc);
     }
 
     // Read the initial parameters (see tes_utils.cpp for more details)
-    read_vector_from_csv(init_param_path_w, weights);
-    read_vector_from_csv(init_param_path_w_sc, weights_sc);
-    read_vector_from_csv(init_param_path_b, bias);
-    read_vector_from_csv(init_param_path_b_sc, bias_sc);
+    read_vector_from_csv(test_saving_paths.init_param_path_w, weights);
+    read_vector_from_csv(test_saving_paths.init_param_path_w_sc, weights_sc);
+    read_vector_from_csv(test_saving_paths.init_param_path_b, bias);
+    read_vector_from_csv(test_saving_paths.init_param_path_b_sc, bias_sc);
 
     // Train the network
     regression_train(tagi_net, train_db, EPOCHS);
@@ -147,18 +130,21 @@ bool test_fnn_full_cov_cpu(bool recompute_outputs, std::string date,
         // RESET OUPUTS
 
         // Write the parameters and hidden states
-        write_vector_to_csv(opt_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(opt_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-        write_vector_to_csv(opt_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(opt_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        write_vector_to_csv(test_saving_paths.opt_param_path_w, "mw,Sw",
+                            weights);
+        write_vector_to_csv(test_saving_paths.opt_param_path_w_sc,
+                            "mw_sc,Sw_sc", weights_sc);
+        write_vector_to_csv(test_saving_paths.opt_param_path_b, "mb,Sb", bias);
+        write_vector_to_csv(test_saving_paths.opt_param_path_b_sc,
+                            "mb_sc,Sb_sc", bias_sc);
 
         // Write the forward hidden states
-        write_vector_to_csv(forward_states_path, "mz,Sz,ma,Sa,J",
-                            forward_states);
+        write_vector_to_csv(test_saving_paths.forward_states_path,
+                            "mz,Sz,ma,Sa,J", forward_states);
 
         // Write the backward hidden states
-        write_vector_to_csv(backward_states_path, backward_states_header,
-                            backward_states_ptr);
+        write_vector_to_csv(test_saving_paths.backward_states_path,
+                            backward_states_header, backward_states_ptr);
 
     } else {
         // PERFORM TESTS
@@ -176,10 +162,12 @@ bool test_fnn_full_cov_cpu(bool recompute_outputs, std::string date,
             ref_bias_sc.push_back(new std::vector<float>());
         }
 
-        read_vector_from_csv(opt_param_path_w, ref_weights);
-        read_vector_from_csv(opt_param_path_w_sc, ref_weights_sc);
-        read_vector_from_csv(opt_param_path_b, ref_bias);
-        read_vector_from_csv(opt_param_path_b_sc, ref_bias_sc);
+        read_vector_from_csv(test_saving_paths.opt_param_path_w, ref_weights);
+        read_vector_from_csv(test_saving_paths.opt_param_path_w_sc,
+                             ref_weights_sc);
+        read_vector_from_csv(test_saving_paths.opt_param_path_b, ref_bias);
+        read_vector_from_csv(test_saving_paths.opt_param_path_b_sc,
+                             ref_bias_sc);
 
         // Compare optimal values with the ones we got
         if (!compare_vectors(ref_weights, weights, data,
@@ -201,7 +189,8 @@ bool test_fnn_full_cov_cpu(bool recompute_outputs, std::string date,
         for (int i = 0; i < 5; i++)
             ref_forward_states.push_back(new std::vector<float>());
 
-        read_vector_from_csv(forward_states_path, ref_forward_states);
+        read_vector_from_csv(test_saving_paths.forward_states_path,
+                             ref_forward_states);
 
         // Compare the saved forward hidden states with the ones we got
         if (!compare_vectors(ref_forward_states, forward_states, data,
@@ -219,7 +208,8 @@ bool test_fnn_full_cov_cpu(bool recompute_outputs, std::string date,
         for (int i = 0; i < 2 * (net.layers.size() - 2); i++)
             ref_backward_states.push_back(new std::vector<float>());
 
-        read_vector_from_csv(backward_states_path, ref_backward_states);
+        read_vector_from_csv(test_saving_paths.backward_states_path,
+                             ref_backward_states);
 
         // Compare the saved backward hidden states with the ones we got
         if (!compare_vectors(ref_backward_states, backward_states_ptr, data,
@@ -233,4 +223,4 @@ bool test_fnn_full_cov_cpu(bool recompute_outputs, std::string date,
         }
     }
     return true;
-}
\ No newline at end of file
+}
diff --git a/test/fnn_heteros/test_fnn_heteros_cpu.cpp b/test/fnn_heteros/test_fnn_heteros_cpu.cpp
index e32781acf..48960703f 100644
--- a/test/fnn_heteros/test_fnn_heteros_cpu.cpp
+++ b/test/fnn_heteros/test_fnn_heteros_cpu.cpp
@@ -47,28 +47,8 @@ bool test_fnn_heteros_cpu(bool recompute_outputs, std::string date,
     path.curr_path = get_current_dir();
     std::string data_path = path.curr_path + "/test/data/" + data;
     std::string data_dir = path.curr_path + "/test/" + arch + "/data/";
-    std::string init_param_path_w =
-        data_dir + date + "_" + data + "_init_param_weights_w.csv";
-    std::string init_param_path_w_sc =
-        data_dir + date + "_" + data + "_init_param_weights_w_sc.csv";
-    std::string init_param_path_b =
-        data_dir + date + "_" + data + "_init_param_bias_b.csv";
-    std::string init_param_path_b_sc =
-        data_dir + date + "_" + data + "_init_param_bias_b_sc.csv";
-    std::string opt_param_path_w =
-        data_dir + date + "_" + data + "_opt_param_weights_w.csv";
-    std::string opt_param_path_w_sc =
-        data_dir + date + "_" + data + "_opt_param_weights_w_sc.csv";
-    std::string opt_param_path_b =
-        data_dir + date + "_" + data + "_opt_param_bias_b.csv";
-    std::string opt_param_path_b_sc =
-        data_dir + date + "_" + data + "_opt_param_bias_b_sc.csv";
-    std::string forward_states_path = data_dir + date +
-                                      "_forward_hidden_states_" + arch + "_" +
-                                      data + ".csv";
-    std::string backward_states_path = data_dir + date +
-                                       "_backward_hidden_states_" + arch + "_" +
-                                       data + ".csv";
+
+    TestSavingPaths test_saving_paths(path.curr_path, arch, data, date);
 
     // Train data
     Dataloader train_db = train_data(data, tagi_net, data_path, NORMALIZE);
@@ -76,19 +56,7 @@ bool test_fnn_heteros_cpu(bool recompute_outputs, std::string date,
     Dataloader test_db =
         test_data(data, tagi_net, data_path, train_db, NORMALIZE);
 
-    std::vector<std::vector<float> *> weights;
-    weights.push_back(&tagi_net.theta.mw);
-    weights.push_back(&tagi_net.theta.Sw);
-    std::vector<std::vector<float> *> weights_sc;
-    weights_sc.push_back(&tagi_net.theta.mw_sc);
-    weights_sc.push_back(&tagi_net.theta.Sw_sc);
-
-    std::vector<std::vector<float> *> bias;
-    bias.push_back(&tagi_net.theta.mb);
-    bias.push_back(&tagi_net.theta.Sb);
-    std::vector<std::vector<float> *> bias_sc;
-    bias_sc.push_back(&tagi_net.theta.mb_sc);
-    bias_sc.push_back(&tagi_net.theta.Sb_sc);
+    TestParamAndStates params_and_states(tagi_net);
 
     // If we want to test but no data is available, we throw an error
     if (!recompute_outputs && !directory_exists(data_dir)) {
@@ -101,91 +69,62 @@ bool test_fnn_heteros_cpu(bool recompute_outputs, std::string date,
             std::cout << "Error: could not create data directory" << std::endl;
             return false;
         }
-        write_vector_to_csv(init_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(init_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-
-        write_vector_to_csv(init_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(init_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, true);
     }
 
     // Read the initial parameters (see tes_utils.cpp for more details)
-    read_vector_from_csv(init_param_path_w, weights);
-    read_vector_from_csv(init_param_path_w_sc, weights_sc);
-    read_vector_from_csv(init_param_path_b, bias);
-    read_vector_from_csv(init_param_path_b_sc, bias_sc);
+    params_and_states.read_params(test_saving_paths, true);
 
     // Train the network
     regression_train(tagi_net, train_db, EPOCHS);
 
-    std::vector<std::vector<float> *> forward_states;
-    forward_states.push_back(&tagi_net.state.mz);
-    forward_states.push_back(&tagi_net.state.Sz);
-    forward_states.push_back(&tagi_net.state.ma);
-    forward_states.push_back(&tagi_net.state.Sa);
-    forward_states.push_back(&tagi_net.state.J);
+    add_forward_states(params_and_states.forward_states, tagi_net);
 
     std::vector<std::vector<float>> backward_states;
     std::string backward_states_header = "";
 
-    for (int i = 0; i < net.layers.size() - 2; i++) {
-        backward_states_header +=
-            "mean_" + std::to_string(i) + ",sigma_" + std::to_string(i) + ",";
-        backward_states.push_back(
-            std::get<0>(tagi_net.get_inovation_mean_var(i)));
-        backward_states.push_back(
-            std::get<1>(tagi_net.get_inovation_mean_var(i)));
-    }
+    add_backward_states(backward_states, backward_states_header, tagi_net,
+                        net.layers.size());
 
-    std::vector<std::vector<float> *> backward_states_ptr;
     for (int i = 0; i < backward_states.size(); i++)
-        backward_states_ptr.push_back(&backward_states[i]);
+        params_and_states.backward_states.push_back(&backward_states[i]);
 
     if (recompute_outputs) {
         // RESET OUPUTS
 
         // Write the parameters and hidden states
-        write_vector_to_csv(opt_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(opt_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-        write_vector_to_csv(opt_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(opt_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, false);
 
         // Write the forward hidden states
-        write_vector_to_csv(forward_states_path, "mz,Sz,ma,Sa,J",
-                            forward_states);
+        write_vector_to_csv(test_saving_paths.forward_states_path,
+                            "mz,Sz,ma,Sa,J", params_and_states.forward_states);
 
         // Write the backward hidden states
-        write_vector_to_csv(backward_states_path, backward_states_header,
-                            backward_states_ptr);
+        write_vector_to_csv(test_saving_paths.backward_states_path,
+                            backward_states_header,
+                            params_and_states.backward_states);
 
     } else {
         // PERFORM TESTS
 
         // Read the saved reference parameters
-        std::vector<std::vector<float> *> ref_weights;
-        std::vector<std::vector<float> *> ref_weights_sc;
-        std::vector<std::vector<float> *> ref_bias;
-        std::vector<std::vector<float> *> ref_bias_sc;
-
-        for (int i = 0; i < 2; i++) {
-            ref_weights.push_back(new std::vector<float>());
-            ref_weights_sc.push_back(new std::vector<float>());
-            ref_bias.push_back(new std::vector<float>());
-            ref_bias_sc.push_back(new std::vector<float>());
-        }
+        TestParamAndStates params_and_states_reference(tagi_net);
 
-        read_vector_from_csv(opt_param_path_w, ref_weights);
-        read_vector_from_csv(opt_param_path_w_sc, ref_weights_sc);
-        read_vector_from_csv(opt_param_path_b, ref_bias);
-        read_vector_from_csv(opt_param_path_b_sc, ref_bias_sc);
+        params_and_states_reference.read_params(test_saving_paths, false);
 
         // Compare optimal values with the ones we got
-        if (!compare_vectors(ref_weights, weights, data,
+        if (!compare_vectors(params_and_states_reference.weights,
+                             params_and_states.weights, data,
                              "fnn heteros weights") ||
-            !compare_vectors(ref_weights_sc, weights_sc, data,
-                             "fnn heteros for residual weights") ||
-            !compare_vectors(ref_bias, bias, data, "fnn heteros bias") ||
-            !compare_vectors(ref_bias_sc, bias_sc, data,
-                             "fnn heteros for residual bias")) {
+            !compare_vectors(params_and_states_reference.weights_sc,
+                             params_and_states.weights_sc, data,
+                             "fnn heteros weights for residual network") ||
+            !compare_vectors(params_and_states_reference.bias,
+                             params_and_states.bias, data,
+                             "fnn heteros bias") ||
+            !compare_vectors(params_and_states_reference.bias_sc,
+                             params_and_states.bias_sc, data,
+                             "fnn heteros bias for residual network")) {
             std::cout
                 << "\033[1;31mTest for FNN HETEROS PARAMS has FAILED in " +
                        data + " data\033[0m\n"
@@ -194,14 +133,16 @@ bool test_fnn_heteros_cpu(bool recompute_outputs, std::string date,
         }
 
         // Read the saved forward hidden states reference
-        std::vector<std::vector<float> *> ref_forward_states;
         for (int i = 0; i < 5; i++)
-            ref_forward_states.push_back(new std::vector<float>());
+            params_and_states_reference.forward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(forward_states_path, ref_forward_states);
+        read_vector_from_csv(test_saving_paths.forward_states_path,
+                             params_and_states_reference.forward_states);
 
         // Compare the saved forward hidden states with the ones we got
-        if (!compare_vectors(ref_forward_states, forward_states, data,
+        if (!compare_vectors(params_and_states_reference.forward_states,
+                             params_and_states.forward_states, data,
                              "fnn heteros forward hidden states")) {
             std::cout
                 << "\033[1;31mTest for FNN HETEROS FORWARD HIDDEN STATES has "
@@ -212,14 +153,16 @@ bool test_fnn_heteros_cpu(bool recompute_outputs, std::string date,
         }
 
         // Read the saved backward hidden states reference
-        std::vector<std::vector<float> *> ref_backward_states;
         for (int i = 0; i < 2 * (net.layers.size() - 2); i++)
-            ref_backward_states.push_back(new std::vector<float>());
+            params_and_states_reference.backward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(backward_states_path, ref_backward_states);
+        read_vector_from_csv(test_saving_paths.backward_states_path,
+                             params_and_states_reference.backward_states);
 
         // Compare the saved backward hidden states with the ones we got
-        if (!compare_vectors(ref_backward_states, backward_states_ptr, data,
+        if (!compare_vectors(params_and_states_reference.backward_states,
+                             params_and_states.backward_states, data,
                              "fnn heteros backward hidden states")) {
             std::cout
                 << "\033[1;31mTest for FNN HETEROS BACKWARD HIDDEN STATES has "
@@ -230,4 +173,4 @@ bool test_fnn_heteros_cpu(bool recompute_outputs, std::string date,
         }
     }
     return true;
-}
\ No newline at end of file
+}
diff --git a/test/lstm/test_lstm_cpu.cpp b/test/lstm/test_lstm_cpu.cpp
index 28a8a2122..5b77de35b 100644
--- a/test/lstm/test_lstm_cpu.cpp
+++ b/test/lstm/test_lstm_cpu.cpp
@@ -54,28 +54,8 @@ bool test_lstm_cpu(bool recompute_outputs, std::string date, std::string arch,
     path.curr_path = get_current_dir();
     std::string data_path = path.curr_path + "/test/data/" + data;
     std::string data_dir = path.curr_path + "/test/" + arch + "/data/";
-    std::string init_param_path_w =
-        data_dir + date + "_" + data + "_init_param_weights_w.csv";
-    std::string init_param_path_w_sc =
-        data_dir + date + "_" + data + "_init_param_weights_w_sc.csv";
-    std::string init_param_path_b =
-        data_dir + date + "_" + data + "_init_param_bias_b.csv";
-    std::string init_param_path_b_sc =
-        data_dir + date + "_" + data + "_init_param_bias_b_sc.csv";
-    std::string opt_param_path_w =
-        data_dir + date + "_" + data + "_opt_param_weights_w.csv";
-    std::string opt_param_path_w_sc =
-        data_dir + date + "_" + data + "_opt_param_weights_w_sc.csv";
-    std::string opt_param_path_b =
-        data_dir + date + "_" + data + "_opt_param_bias_b.csv";
-    std::string opt_param_path_b_sc =
-        data_dir + date + "_" + data + "_opt_param_bias_b_sc.csv";
-    std::string forward_states_path = data_dir + date +
-                                      "_forward_hidden_states_" + arch + "_" +
-                                      data + ".csv";
-    std::string backward_states_path = data_dir + date +
-                                       "_backward_hidden_states_" + arch + "_" +
-                                       data + ".csv";
+
+    TestSavingPaths test_saving_paths(path.curr_path, arch, data, date);
 
     // Train data
     Dataloader train_db = test_time_series_datloader(
@@ -85,19 +65,7 @@ bool test_lstm_cpu(bool recompute_outputs, std::string date, std::string arch,
     Dataloader test_db = test_time_series_datloader(
         tagi_net.prop, "test", NUM_FEATURES, data_path, OUTPUT_COL, NORMALIZE);
 
-    std::vector<std::vector<float> *> weights;
-    weights.push_back(&tagi_net.theta.mw);
-    weights.push_back(&tagi_net.theta.Sw);
-    std::vector<std::vector<float> *> weights_sc;
-    weights_sc.push_back(&tagi_net.theta.mw_sc);
-    weights_sc.push_back(&tagi_net.theta.Sw_sc);
-
-    std::vector<std::vector<float> *> bias;
-    bias.push_back(&tagi_net.theta.mb);
-    bias.push_back(&tagi_net.theta.Sb);
-    std::vector<std::vector<float> *> bias_sc;
-    bias_sc.push_back(&tagi_net.theta.mb_sc);
-    bias_sc.push_back(&tagi_net.theta.Sb_sc);
+    TestParamAndStates params_and_states(tagi_net);
 
     // If we want to test but no data is available, we throw an error
     if (!recompute_outputs && !directory_exists(data_dir)) {
@@ -110,89 +78,59 @@ bool test_lstm_cpu(bool recompute_outputs, std::string date, std::string arch,
             std::cout << "Error: could not create data directory" << std::endl;
             return false;
         }
-        write_vector_to_csv(init_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(init_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-
-        write_vector_to_csv(init_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(init_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, true);
     }
 
     // Read the initial parameters (see tes_utils.cpp for more details)
-    read_vector_from_csv(init_param_path_w, weights);
-    read_vector_from_csv(init_param_path_w_sc, weights_sc);
-    read_vector_from_csv(init_param_path_b, bias);
-    read_vector_from_csv(init_param_path_b_sc, bias_sc);
+    params_and_states.read_params(test_saving_paths, true);
 
     // Train the network
     train_time_series(tagi_net, train_db, EPOCHS);
 
-    std::vector<std::vector<float> *> forward_states;
-    forward_states.push_back(&tagi_net.state.mz);
-    forward_states.push_back(&tagi_net.state.Sz);
-    forward_states.push_back(&tagi_net.state.ma);
-    forward_states.push_back(&tagi_net.state.Sa);
-    forward_states.push_back(&tagi_net.state.J);
+    add_forward_states(params_and_states.forward_states, tagi_net);
 
     std::vector<std::vector<float>> backward_states;
     std::string backward_states_header = "";
 
-    for (int i = 0; i < net.layers.size() - 2; i++) {
-        backward_states_header +=
-            "mean_" + std::to_string(i) + ",sigma_" + std::to_string(i) + ",";
-        backward_states.push_back(
-            std::get<0>(tagi_net.get_inovation_mean_var(i)));
-        backward_states.push_back(
-            std::get<1>(tagi_net.get_inovation_mean_var(i)));
-    }
+    add_backward_states(backward_states, backward_states_header, tagi_net,
+                        net.layers.size());
 
-    std::vector<std::vector<float> *> backward_states_ptr;
     for (int i = 0; i < backward_states.size(); i++)
-        backward_states_ptr.push_back(&backward_states[i]);
+        params_and_states.backward_states.push_back(&backward_states[i]);
 
     if (recompute_outputs) {
         // RESET OUPUTS
 
         // Write the parameters and hidden states
-        write_vector_to_csv(opt_param_path_w, "mw,Sw", weights);
-        write_vector_to_csv(opt_param_path_w_sc, "mw_sc,Sw_sc", weights_sc);
-        write_vector_to_csv(opt_param_path_b, "mb,Sb", bias);
-        write_vector_to_csv(opt_param_path_b_sc, "mb_sc,Sb_sc", bias_sc);
+        params_and_states.write_params(test_saving_paths, false);
 
         // Write the forward hidden states
-        write_vector_to_csv(forward_states_path, "mz,Sz,ma,Sa,J",
-                            forward_states);
+        write_vector_to_csv(test_saving_paths.forward_states_path,
+                            "mz,Sz,ma,Sa,J", params_and_states.forward_states);
 
         // Write the backward hidden states
-        write_vector_to_csv(backward_states_path, backward_states_header,
-                            backward_states_ptr);
+        write_vector_to_csv(test_saving_paths.backward_states_path,
+                            backward_states_header,
+                            params_and_states.backward_states);
 
     } else {
         // PERFORM TESTS
 
         // Read the saved reference parameters
-        std::vector<std::vector<float> *> ref_weights;
-        std::vector<std::vector<float> *> ref_weights_sc;
-        std::vector<std::vector<float> *> ref_bias;
-        std::vector<std::vector<float> *> ref_bias_sc;
-
-        for (int i = 0; i < 2; i++) {
-            ref_weights.push_back(new std::vector<float>());
-            ref_weights_sc.push_back(new std::vector<float>());
-            ref_bias.push_back(new std::vector<float>());
-            ref_bias_sc.push_back(new std::vector<float>());
-        }
+        TestParamAndStates params_and_states_reference(tagi_net);
 
-        read_vector_from_csv(opt_param_path_w, ref_weights);
-        read_vector_from_csv(opt_param_path_w_sc, ref_weights_sc);
-        read_vector_from_csv(opt_param_path_b, ref_bias);
-        read_vector_from_csv(opt_param_path_b_sc, ref_bias_sc);
+        params_and_states_reference.read_params(test_saving_paths, false);
 
         // Compare optimal values with the ones we got
-        if (!compare_vectors(ref_weights, weights, data, "lstm weights") ||
-            !compare_vectors(ref_weights_sc, weights_sc, data,
+        if (!compare_vectors(params_and_states_reference.weights,
+                             params_and_states.weights, data, "lstm weights") ||
+            !compare_vectors(params_and_states_reference.weights_sc,
+                             params_and_states.weights_sc, data,
                              "lstm weights for residual network") ||
-            !compare_vectors(ref_bias, bias, data, "lstm bias") ||
-            !compare_vectors(ref_bias_sc, bias_sc, data,
+            !compare_vectors(params_and_states_reference.bias,
+                             params_and_states.bias, data, "lstm bias") ||
+            !compare_vectors(params_and_states_reference.bias_sc,
+                             params_and_states.bias_sc, data,
                              "lstm bias for residual network")) {
             std::cout << "\033[1;31mTest for LSTM PARAMS has FAILED in " +
                              data + " data\033[0m\n"
@@ -201,14 +139,16 @@ bool test_lstm_cpu(bool recompute_outputs, std::string date, std::string arch,
         }
 
         // Read the saved forward hidden states reference
-        std::vector<std::vector<float> *> ref_forward_states;
         for (int i = 0; i < 5; i++)
-            ref_forward_states.push_back(new std::vector<float>());
+            params_and_states_reference.forward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(forward_states_path, ref_forward_states);
+        read_vector_from_csv(test_saving_paths.forward_states_path,
+                             params_and_states_reference.forward_states);
 
         // Compare the saved forward hidden states with the ones we got
-        if (!compare_vectors(ref_forward_states, forward_states, data,
+        if (!compare_vectors(params_and_states_reference.forward_states,
+                             params_and_states.forward_states, data,
                              "lstm forward hidden states")) {
             std::cout << "\033[1;31mTest for LSTM FORWARD HIDDEN STATES has "
                          "FAILED in " +
@@ -218,14 +158,16 @@ bool test_lstm_cpu(bool recompute_outputs, std::string date, std::string arch,
         }
 
         // Read the saved backward hidden states reference
-        std::vector<std::vector<float> *> ref_backward_states;
         for (int i = 0; i < 2 * (net.layers.size() - 2); i++)
-            ref_backward_states.push_back(new std::vector<float>());
+            params_and_states_reference.backward_states.push_back(
+                new std::vector<float>());
 
-        read_vector_from_csv(backward_states_path, ref_backward_states);
+        read_vector_from_csv(test_saving_paths.backward_states_path,
+                             params_and_states_reference.backward_states);
 
         // Compare the saved backward hidden states with the ones we got
-        if (!compare_vectors(ref_backward_states, backward_states_ptr, data,
+        if (!compare_vectors(params_and_states_reference.backward_states,
+                             params_and_states.backward_states, data,
                              "lstm backward hidden states")) {
             std::cout << "\033[1;31mTest for LSTM BACKWARD HIDDEN STATES has "
                          "FAILED in " +
@@ -235,4 +177,4 @@ bool test_lstm_cpu(bool recompute_outputs, std::string date, std::string arch,
         }
     }
     return true;
-}
\ No newline at end of file
+}
diff --git a/test/test_autoencoder.cu b/test/test_autoencoder.cu
new file mode 100644
index 000000000..667200454
--- /dev/null
+++ b/test/test_autoencoder.cu
@@ -0,0 +1,84 @@
+///////////////////////////////////////////////////////////////////////////////
+// File:         test_autoencoder.cu
+// Description:  Auxiliar independent script to perform image generation
+// Authors:      Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet
+// Created:      April 12, 2023
+// Updated:      April 13, 2023
+// Contact:      miquelflorensa11@gmail.com & luongha.nguyen@gmail.com &
+//               james.goulet@polymtl.ca
+// Copyright (c) 2023 Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet.
+// Some rights reserved.
+///////////////////////////////////////////////////////////////////////////////
+
+#include "test_autoencoder.cuh"
+
+void train_autoencoder(TagiNetwork &net_e, TagiNetwork &net_d, ImageData &imdb,
+                       int n_classes) {
+    // Batch size check
+    if (net_e.prop.batch_size != net_d.prop.batch_size) {
+        throw std::invalid_argument(
+            " Batch size is not equal - Task - Autoencoder");
+    }
+
+    // Compute number of data
+    int n_iter = imdb.num_data / net_d.prop.batch_size;
+
+    // Input and output layer
+    std::vector<float> x_batch, Sx_batch, y_batch, y_batch_e, V_batch_e,
+        delta_mz, delta_Sz;
+    std::vector<int> data_idx = create_range(imdb.num_data);
+    std::vector<int> test_data_idx = create_range(imdb.num_data);
+    std::vector<int> batch_idx(net_d.prop.batch_size);
+    std::vector<int> idx_ud_batch(net_d.prop.nye * net_d.prop.batch_size, 0);
+    std::vector<int> idx_ud_batch_e(net_e.prop.nye * net_e.prop.batch_size, 0);
+    std::vector<int> label_batch(net_d.prop.batch_size, 0);
+
+    x_batch.resize(net_e.prop.batch_size * net_e.prop.n_x, 0);
+    Sx_batch.resize(net_e.prop.batch_size * net_e.prop.n_x,
+                    powf(net_e.prop.sigma_x, 2));
+    y_batch.resize(net_d.prop.batch_size * net_d.prop.n_y, 0);
+    y_batch_e.resize(net_e.prop.batch_size * net_e.prop.n_y, 0);
+    V_batch_e.resize(net_e.prop.batch_size * net_e.prop.n_y, 0);
+
+    std::vector<float> Sx_f_batch;
+
+    std::vector<float> V_batch(net_d.prop.batch_size * net_d.prop.n_y,
+                               powf(net_d.prop.sigma_v, 2));
+
+    for (int i = 0; i < n_iter; i++) {
+        ;
+        // TODO: Make a cleaner way to handle both cases
+        if (i == 0) {
+            net_e.prop.ra_mt = 0.0f;
+            net_d.prop.ra_mt = 0.0f;
+        } else {
+            net_e.prop.ra_mt = 0.9f;
+            net_d.prop.ra_mt = 0.9f;
+        }
+
+        // Load input data for encoder and output data for decoder
+        get_batch_idx(data_idx, i * net_d.prop.batch_size,
+                      net_e.prop.batch_size, batch_idx);
+        get_batch_data(imdb.images, batch_idx, net_e.prop.n_x, x_batch);
+        get_batch_data(imdb.labels, batch_idx, 1, label_batch);
+
+        // Feed forward for encoder
+        net_e.feed_forward(x_batch, Sx_batch, Sx_f_batch);
+
+        // Get all output's states of encoder
+        net_e.get_all_network_outputs();
+
+        // Feed forward for decoder
+        net_d.connected_feed_forward(net_e.ma, net_e.Sa, net_e.mz, net_e.Sz,
+                                     net_e.J);
+
+        // Feed backward for decoder
+        net_d.state_feed_backward(x_batch, V_batch, idx_ud_batch);
+        net_d.param_feed_backward();
+
+        // Feed backward for encoder
+        std::tie(delta_mz, delta_Sz) = net_d.get_state_delta_mean_var();
+        net_e.state_feed_backward(delta_mz, delta_Sz, idx_ud_batch);
+        net_e.param_feed_backward();
+    }
+}
diff --git a/test/test_autoencoder.cuh b/test/test_autoencoder.cuh
new file mode 100644
index 000000000..455945292
--- /dev/null
+++ b/test/test_autoencoder.cuh
@@ -0,0 +1,44 @@
+///////////////////////////////////////////////////////////////////////////////
+// File:         test_autoencoder.cuh
+// Description:  Header of script to perform independent image generation
+// Authors:      Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet
+// Created:      April 12, 2023
+// Updated:      April 13, 2023
+// Contact:      miquelflorensa11@gmail.com & luongha.nguyen@gmail.com &
+//               james.goulet@polymtl.ca
+// Copyright (c) 2023 Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet.
+// Some rights reserved.
+///////////////////////////////////////////////////////////////////////////////
+
+#pragma once
+#include <iomanip>
+#include <iostream>
+#include <string>
+
+#include "../include/common.h"
+#include "../include/cost.h"
+#include "../include/data_transfer.cuh"
+#include "../include/dataloader.h"
+#include "../include/feed_forward.cuh"
+#include "../include/global_param_update.cuh"
+#include "../include/gpu_debug_utils.h"
+#include "../include/indices.h"
+#include "../include/lstm_feed_forward.cuh"
+#include "../include/net_init.h"
+#include "../include/net_prop.h"
+#include "../include/param_feed_backward.cuh"
+#include "../include/state_feed_backward.cuh"
+#include "../include/struct_var.h"
+#include "../include/tagi_network.cuh"
+#include "../include/user_input.h"
+
+/**
+ * @brief Auxiliar autoencoder function for unit tests
+ *
+ * @param net_e  The encoder network
+ * @param net_d  The decoder network
+ * @param imdb The image train data
+ * @param n_classes The number of classes
+ */
+void train_autoencoder(TagiNetwork &net_e, TagiNetwork &net_d, ImageData &imdb,
+                       int n_classes);
diff --git a/test/test_cpu.cpp b/test/test_cpu.cpp
index 9fa3e1580..ec0f538bc 100644
--- a/test/test_cpu.cpp
+++ b/test/test_cpu.cpp
@@ -3,7 +3,7 @@
 // Description:  Main script to test the CPU implementation of cuTAGI
 // Authors:      Florensa, Miquel, Luong-Ha Nguyen & James-A. Goulet
 // Created:      February 20, 2023
-// Updated:      April 4, 2023
+// Updated:      April 13, 2023
 // Contact:      miquelflorensa11@gmail.com, luongha.nguyen@gmail.com &
 //               james.goulet@polymtl.ca
 // Copyright (c) 2023 Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet.
@@ -15,8 +15,8 @@
 const int NUM_TESTS_CPU = 6;
 const int NUM_TESTS_GPU = 3;
 const std::vector<std::string> AVAILABLE_ARCHITECTURES = {
-    "all",  "fnn",      "fnn_heteros", "fnn_full_cov",   "fnn_derivatives",
-    "lstm", "act_func", "cnn",         "cnn_batch_norm", "autoencoder"};
+    "fnn",      "fnn_heteros", "fnn_full_cov",   "fnn_derivatives", "lstm",
+    "act_func", "cnn",         "cnn_batch_norm", "autoencoder"};
 
 std::vector<std::string> read_dates() {
     std::ifstream file("test/data/last_dates.csv");
@@ -43,7 +43,7 @@ void write_dates(std::vector<std::string> dates, int column, std::string date) {
         return;
     }
 
-    for (int i = 1; i < AVAILABLE_ARCHITECTURES.size(); i++) {
+    for (int i = 0; i < AVAILABLE_ARCHITECTURES.size(); i++) {
         file << AVAILABLE_ARCHITECTURES[i];
         if (i != AVAILABLE_ARCHITECTURES.size() - 1) {
             file << ",";
@@ -78,27 +78,32 @@ void check_valid_input_architecture(std::string test_architecture) {
 }
 
 void print_test_results(bool single_test, bool test_passed, int num_tests,
-                        int test_num, std::string arch_name) {
-    std::string output_color;
-    std::string output;
+                        int test_num, std::string arch_name,
+                        std::chrono::milliseconds run_time) {
+    if (!single_test) {
+        int percentage = round((100.0 / num_tests) * (test_num + 1));
+        std::cout << "[";
+        if (percentage < 100) std::cout << " ";
+        if (percentage < 10) std::cout << " ";
+        std::cout << percentage << "%] ";
+    }
+    int num_spaces = 35;
+    arch_name.append(":");
 
     if (test_passed) {
-        output_color = "[32;1m";
-        output = "passed";
+        std::cout << std::setw(num_spaces) << std::left << arch_name;
+        std::cout << "\033[32;1mPASS\033[0m";
     } else {
-        output_color = "[31;1m";
-        output = "failed";
+        std::cout << std::setw(num_spaces) << std::left << arch_name;
+        std::cout << "\033[31;1mFAIL\033[0m";
     }
 
-    if (!single_test) {
-        std::cout << "[ " << floor((100 / num_tests) * (test_num + 1)) << "%] ";
-    }
-    std::cout << "\033" << output_color << arch_name << " tests " << output
-              << "\033[0m" << std::endl;
+    std::cout << "  (took " << run_time.count() << "ms)" << std::endl;
 }
 
 int test_cpu(std::vector<std::string>& user_input_options,
-             bool compute_gpu_tests) {
+             bool compute_gpu_tests,
+             std::chrono::steady_clock::time_point test_start) {
     std::string reinizialize_test_outputs = "";
     std::string test_architecture = "";
     std::string date = "";
@@ -115,10 +120,10 @@ int test_cpu(std::vector<std::string>& user_input_options,
                   << "Perform tests on all architectures" << std::endl;
 
         std::cout << std::setw(num_spaces) << std::left
-                  << "test [architecture-name]"
+                  << "test <architecture-name>"
                   << "Run one specific test" << std::endl;
 
-        std::cout << std::setw(num_spaces) << std::left << "test -reset all"
+        std::cout << std::setw(num_spaces) << std::left << "test -reset (all)"
                   << "Reinizialize all test references" << std::endl;
 
         std::cout << std::setw(num_spaces) << std::left
@@ -140,6 +145,9 @@ int test_cpu(std::vector<std::string>& user_input_options,
         if (user_input_options[0] == "-reset") {
             if (user_input_options.size() == 1) {
                 reinizialize_test_outputs = "all";
+            } else if (user_input_options.size() == 2 &&
+                       user_input_options[1] == "all") {
+                reinizialize_test_outputs = "all";
             } else {
                 // Check if the architecture is valid
                 check_valid_input_architecture(user_input_options[1]);
@@ -192,6 +200,10 @@ int test_cpu(std::vector<std::string>& user_input_options,
             device = "GPU";
         }
 
+        std::cout << "=========================================" << std::endl;
+        std::cout << "            TAGI Unit Test               " << std::endl;
+        std::cout << "=========================================" << std::endl;
+
         int num_test_passed = 0;
 
         std::cout << "Performing " << test_architecture << " tests on "
@@ -201,13 +213,18 @@ int test_cpu(std::vector<std::string>& user_input_options,
         if (test_architecture == "all" || test_architecture == "fnn") {
             test_num = 0;  // FNN
 
+            auto start = std::chrono::steady_clock::now();
             bool test_result =
                 test_fnn_cpu(false, test_dates[test_num], "fnn", "1D") &&
                 test_fnn_cpu(false, test_dates[test_num], "fnn",
                              "Boston_housing");
 
+            auto end = std::chrono::steady_clock::now();
+            auto run_time =
+                std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                      start);
             print_test_results(single_test, test_result, num_tests, test_num,
-                               "FNN");
+                               "FNN", run_time);
 
             if (test_result) num_test_passed++;
         }
@@ -217,12 +234,17 @@ int test_cpu(std::vector<std::string>& user_input_options,
         if (test_architecture == "all" || test_architecture == "fnn_heteros") {
             test_num = 1;  // FNN heteroscedastic noise
 
+            auto start = std::chrono::steady_clock::now();
             bool test_result =
                 test_fnn_heteros_cpu(false, test_dates[test_num], "fnn_heteros",
                                      "1D_noise_inferance");
 
+            auto end = std::chrono::steady_clock::now();
+            auto run_time =
+                std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                      start);
             print_test_results(single_test, test_result, num_tests, test_num,
-                               "FNN heteroscedastic");
+                               "FNN heteroscedastic", run_time);
 
             if (test_result) num_test_passed++;
         }
@@ -231,11 +253,16 @@ int test_cpu(std::vector<std::string>& user_input_options,
         if (test_architecture == "all" || test_architecture == "fnn_full_cov") {
             test_num = 2;  // FNN full covariance
 
+            auto start = std::chrono::steady_clock::now();
             bool test_result = test_fnn_full_cov_cpu(
                 false, test_dates[test_num], "fnn_full_cov", "1D_full_cov");
 
+            auto end = std::chrono::steady_clock::now();
+            auto run_time =
+                std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                      start);
             print_test_results(single_test, test_result, num_tests, test_num,
-                               "FNN full covariance");
+                               "FNN full covariance", run_time);
 
             if (test_result) num_test_passed++;
         }
@@ -246,12 +273,17 @@ int test_cpu(std::vector<std::string>& user_input_options,
             test_architecture == "fnn_derivatives") {
             test_num = 3;  // FNN derivatives
 
+            auto start = std::chrono::steady_clock::now();
             bool test_result =
                 test_fnn_derivatives_cpu(false, test_dates[test_num],
                                          "fnn_derivatives", "1D_derivatives");
 
+            auto end = std::chrono::steady_clock::now();
+            auto run_time =
+                std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                      start);
             print_test_results(single_test, test_result, num_tests, test_num,
-                               "FNN derivatives");
+                               "FNN derivatives", run_time);
 
             if (test_result) num_test_passed++;
         }
@@ -260,11 +292,16 @@ int test_cpu(std::vector<std::string>& user_input_options,
         if (test_architecture == "all" || test_architecture == "lstm") {
             test_num = 4;  // LSTM
 
+            auto start = std::chrono::steady_clock::now();
             bool test_result = test_lstm_cpu(false, test_dates[test_num],
                                              "lstm", "time_series");
 
+            auto end = std::chrono::steady_clock::now();
+            auto run_time =
+                std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                      start);
             print_test_results(single_test, test_result, num_tests, test_num,
-                               "LSTM");
+                               "LSTM", run_time);
 
             if (test_result) num_test_passed++;
         }
@@ -273,22 +310,36 @@ int test_cpu(std::vector<std::string>& user_input_options,
         if (test_architecture == "all" || test_architecture == "act_func") {
             test_num = 5;  // Activation Functions
 
+            auto start = std::chrono::steady_clock::now();
             bool test_result = test_act_func_cpu(false, test_dates[test_num],
                                                  "act_func", "Boston_housing");
 
+            auto end = std::chrono::steady_clock::now();
+            auto run_time =
+                std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                      start);
             print_test_results(single_test, test_result, num_tests, test_num,
-                               "Activation functions");
+                               "Activation functions", run_time);
 
             if (test_result) num_test_passed++;
         }
 
+        auto end = std::chrono::steady_clock::now();
+        auto run_time = std::chrono::duration_cast<std::chrono::milliseconds>(
+                            end - test_start)
+                            .count();
+
         // Number of tests passed
         if (test_architecture == "all" && !compute_gpu_tests) {
             std::cout << std::endl;
-            std::cout << "--------------------SUMMARY--------------------"
+            std::cout << "---------------SUMMARY-------------------"
+                      << std::endl;
+            std::cout << "Total tests: " << num_tests << std::endl;
+            std::cout << "Passed: " << num_test_passed << std::endl;
+            std::cout << "Failed: " << num_tests - num_test_passed << std::endl;
+            std::cout << "Total time taken: " << run_time << "ms" << std::endl;
+            std::cout << "========================================="
                       << std::endl;
-            std::cout << "Passed tests: [" << num_test_passed << "/"
-                      << num_tests << "]" << std::endl;
             return num_test_passed;
         }
         return num_test_passed;
diff --git a/test/test_cpu.h b/test/test_cpu.h
index e9703006b..adfaca84a 100644
--- a/test/test_cpu.h
+++ b/test/test_cpu.h
@@ -54,10 +54,12 @@ void check_valid_input_architecture(std::string test_architecture);
  *
  * @param user_input_options vector with the user input options
  * @param compute_gpu_tests true if the gpu tests will be also computed
+ * @param test_start time when the test started
  * @return Returns number of passed test or -1 for any error
  */
 int test_cpu(std::vector<std::string> &user_input_options,
-             bool compute_gpu_tests);
+             bool compute_gpu_tests,
+             std::chrono::steady_clock::time_point test_start);
 
 /**
  * @brief Print the test results
@@ -66,6 +68,8 @@ int test_cpu(std::vector<std::string> &user_input_options,
  * @param num_tests number of tests
  * @param test_num number of the test
  * @param arch_name name of the architecture
+ * @param run_time time taken to run the test
  */
 void print_test_results(bool single_test, bool test_passed, int num_tests,
-                        int test_num, std::string arch_name);
+                        int test_num, std::string arch_name,
+                        std::chrono::milliseconds run_time);
diff --git a/test/test_gpu.cu b/test/test_gpu.cu
index c1ba12692..499735223 100644
--- a/test/test_gpu.cu
+++ b/test/test_gpu.cu
@@ -3,7 +3,7 @@
 // Description:  Main script to test the GPU implementation of cuTAGI
 // Authors:      Florensa, Miquel, Luong-Ha Nguyen & James-A. Goulet
 // Created:      February 20, 2023
-// Updated:      April 5, 2023
+// Updated:      April 13, 2023
 // Contact:      miquelflorensa11@gmail.com, luongha.nguyen@gmail.com &
 //               james.goulet@polymtl.ca
 // Copyright (c) 2023 Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet.
@@ -15,7 +15,8 @@
 const int NUM_TESTS = 9;
 
 int test_gpu(std::vector<std::string>& user_input_options,
-             int num_tests_passed_cpu) {
+             int num_tests_passed_cpu,
+             std::chrono::steady_clock::time_point test_start) {
     std::string reinizialize_test_outputs = "";
     std::string test_architecture = "";
     std::string date = "";
@@ -23,12 +24,15 @@ int test_gpu(std::vector<std::string>& user_input_options,
 
     if (user_input_options.size() == 1 &&
         (user_input_options[0] == "-h" || user_input_options[0] == "--help")) {
-        // Help message have alrady been showed in test_cpu.cpp
+        // Help message have already been showed in test_cpu.cpp
         return -1;
     } else if (user_input_options.size() > 0 && user_input_options.size() < 3) {
         if (user_input_options[0] == "-reset") {
             if (user_input_options.size() == 1) {
                 reinizialize_test_outputs = "all";
+            } else if (user_input_options.size() == 2 &&
+                       user_input_options[1] == "all") {
+                reinizialize_test_outputs = "all";
             } else {
                 // Check if the architecture is valid
                 check_valid_input_architecture(user_input_options[1]);
@@ -76,11 +80,16 @@ int test_gpu(std::vector<std::string>& user_input_options,
         if (test_architecture == "all" || test_architecture == "cnn") {
             test_num = 6;  // CNN
 
+            auto start = std::chrono::steady_clock::now();
             bool test_result =
                 test_cnn_gpu(false, test_dates[test_num], "cnn", "mnist");
 
+            auto end = std::chrono::steady_clock::now();
+            auto run_time =
+                std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                      start);
             print_test_results(single_test, test_result, NUM_TESTS, test_num,
-                               "CNN");
+                               "CNN", run_time);
 
             if (test_result) num_test_passed++;
         }
@@ -90,22 +99,54 @@ int test_gpu(std::vector<std::string>& user_input_options,
             test_architecture == "cnn_batch_norm") {
             test_num = 7;  // CNN batch norm.
 
+            auto start = std::chrono::steady_clock::now();
             bool test_result = test_cnn_batch_norm_gpu(
                 false, test_dates[test_num], "cnn_batch_norm", "mnist");
 
+            auto end = std::chrono::steady_clock::now();
+            auto run_time =
+                std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                      start);
+            print_test_results(single_test, test_result, NUM_TESTS, test_num,
+                               "CNN batch normalization", run_time);
+
+            if (test_result) num_test_passed++;
+        }
+
+        // Perform test on GPU with autoencoder for image generation
+        if (test_architecture == "all" || test_architecture == "autoencoder") {
+            test_num = 8;  // Autoencoder
+
+            auto start = std::chrono::steady_clock::now();
+            bool test_result = test_autoencoder_gpu(false, test_dates[test_num],
+                                                    "autoencoder", "mnist");
+
+            auto end = std::chrono::steady_clock::now();
+            auto run_time =
+                std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                      start);
             print_test_results(single_test, test_result, NUM_TESTS, test_num,
-                               "CNN batch normalization");
+                               "Autoencoder", run_time);
 
             if (test_result) num_test_passed++;
         }
 
+        auto end = std::chrono::steady_clock::now();
+        auto run_time = std::chrono::duration_cast<std::chrono::milliseconds>(
+                            end - test_start)
+                            .count();
+
         // Number of tests passed
         if (test_architecture == "all") {
             std::cout << std::endl;
-            std::cout << "--------------------SUMMARY--------------------"
+            std::cout << "---------------SUMMARY-------------------"
+                      << std::endl;
+            std::cout << "Total tests: " << NUM_TESTS << std::endl;
+            std::cout << "Passed: " << num_test_passed << std::endl;
+            std::cout << "Failed: " << NUM_TESTS - num_test_passed << std::endl;
+            std::cout << "Total time taken: " << run_time << "ms" << std::endl;
+            std::cout << "========================================="
                       << std::endl;
-            std::cout << "Passed tests: [" << num_test_passed << "/"
-                      << NUM_TESTS << "]" << std::endl;
             return num_test_passed;
         }
         return -1;
@@ -144,6 +185,20 @@ int test_gpu(std::vector<std::string>& user_input_options,
             write_dates(test_dates, test_num, date);
             test_dates[test_num] = date;
         }
+
+        if (reinizialize_test_outputs == "all" ||
+            reinizialize_test_outputs == "autoencoder") {
+            // Reinizialize test outputs for for image generation
+            std::cout << "Reinizializing Autoencoder test outputs" << std::endl;
+
+            test_autoencoder_gpu(true, date, "autoencoder", "mnist");
+
+            test_num = 8;  // Autoencoder
+
+            // Update de last date of the test
+            write_dates(test_dates, test_num, date);
+            test_dates[test_num] = date;
+        }
         return 0;
     }
 
diff --git a/test/test_gpu.cuh b/test/test_gpu.cuh
index a9e9d1f54..7fe24cd40 100644
--- a/test/test_gpu.cuh
+++ b/test/test_gpu.cuh
@@ -4,7 +4,7 @@
 //               of cuTAGI
 // Authors:      Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet
 // Created:      February 20, 2023
-// Updated:      March 21, 2023
+// Updated:      April 13, 2023
 // Contact:      miquelflorensa11@gmail.com & luongha.nguyen@gmail.com &
 //               james.goulet@polymtl.ca
 // Copyright (c) 2023 Miquel Florensa, Luong-Ha Nguyen & James-A. Goulet.
@@ -19,6 +19,7 @@
 #include <sstream>
 #include <string>
 
+#include "autoencoder/test_autoencoder_gpu.cuh"
 #include "cnn/test_cnn_gpu.cuh"
 #include "cnn_batch_norm/test_cnn_batch_norm_gpu.cuh"
 #include "test_cpu.h"
@@ -28,7 +29,9 @@
  *
  * @param user_input_options vector with the user input options
  * @param num_tests_passed_cpu number of cpu passed
+ * @param test_start time point of the start of the test
  * @return Returns number of passed test or -1 for any error
  */
 int test_gpu(std::vector<std::string> &user_input_options,
-             int num_tests_passed_cpu);
+             int num_tests_passed_cpu,
+             std::chrono::steady_clock::time_point test_start);
diff --git a/test/test_utils.h b/test/test_utils.h
index c8cf5b2ea..cc57f17c2 100644
--- a/test/test_utils.h
+++ b/test/test_utils.h
@@ -3,7 +3,7 @@
 // Description:  Header file for the utils functions for unitest
 // Authors:      Florensa, Miquel & Luong-Ha Nguyen & James-A. Goulet
 // Created:      February 20, 2023
-// Updated:      February 20, 2023
+// Updated:      April 13, 2023
 // Contact:      miquelflorensa11@gmail.com & luongha.nguyen@gmail.com &
 //               james.goulet@polymtl.ca
 // Copyright (c) 2022 Luong-Ha Nguyen & James-A. Goulet.
@@ -40,6 +40,68 @@
 #include "../include/net_prop.h"
 #include "../include/struct_var.h"
 
+/**
+ * @brief Initialize the weights and bias
+ *
+ * @param weights weights
+ * @param weights_sc weights standard deviation
+ * @param bias bias
+ * @param bias_sc bias standard deviation
+ * @param net neural network
+ */
+template <typename Net>
+void add_weights_and_bias(std::vector<std::vector<float> *> &weights,
+                          std::vector<std::vector<float> *> &weights_sc,
+                          std::vector<std::vector<float> *> &bias,
+                          std::vector<std::vector<float> *> &bias_sc,
+                          Net &net) {
+    weights.push_back(&net.theta.mw);
+    weights.push_back(&net.theta.Sw);
+
+    weights_sc.push_back(&net.theta.mw_sc);
+    weights_sc.push_back(&net.theta.Sw_sc);
+
+    bias.push_back(&net.theta.mb);
+    bias.push_back(&net.theta.Sb);
+
+    bias_sc.push_back(&net.theta.mb_sc);
+    bias_sc.push_back(&net.theta.Sb_sc);
+}
+
+/**
+ * @brief Initialize the forward states
+ *
+ * @param forward_states forward states
+ * @param net neural network
+ */
+template <typename Net>
+void add_forward_states(std::vector<std::vector<float> *> &forward_states,
+                        Net &net) {
+    forward_states.push_back(&net.state.mz);
+    forward_states.push_back(&net.state.Sz);
+    forward_states.push_back(&net.state.ma);
+    forward_states.push_back(&net.state.Sa);
+    forward_states.push_back(&net.state.J);
+}
+
+/**
+ * @brief Initialize the backward states
+ *
+ * @param backward_states backward states
+ * @param net neural network
+ */
+template <typename Net>
+void add_backward_states(std::vector<std::vector<float>> &backward_states,
+                         std::string &backward_states_header, Net &net,
+                         int layers) {
+    for (int i = 0; i < layers - 2; i++) {
+        backward_states_header +=
+            "mean_" + std::to_string(i) + ",sigma_" + std::to_string(i) + ",";
+        backward_states.push_back(std::get<0>(net.get_inovation_mean_var(i)));
+        backward_states.push_back(std::get<1>(net.get_inovation_mean_var(i)));
+    }
+}
+
 /**
  * @brief Indicate if a directory exists
  *
@@ -201,4 +263,130 @@ void write_vector_to_csv(std::string filename, std::string header,
     }
 
     file.close();
-}
\ No newline at end of file
+}
+
+/**
+ * @brief Class to store the paths to the data files
+ */
+class TestSavingPaths {
+   public:
+    /**
+     * @brief Construct a new Test Saving Paths object
+     *
+     * @param curr_path current path
+     * @param arch architecture
+     * @param data data
+     * @param date date
+     * @param add_encoder add encoder
+     * @param add_decoder add decoder
+     */
+    TestSavingPaths(std::string curr_path, std::string arch, std::string data,
+                    std::string date, bool add_encoder = false,
+                    bool add_decoder = false) {
+        std::string encoder_suffix = add_encoder ? "encoder_" : "";
+        std::string decoder_suffix = add_decoder ? "decoder_" : "";
+        std::string data_dir = curr_path + "/test/" + arch + "/data/" + date +
+                               "_" + encoder_suffix + decoder_suffix;
+        std::string path_sufix = "_" + arch + "_" + data + ".csv";
+
+        init_param_path_w = data_dir + "init_param_weights_w" + path_sufix;
+        init_param_path_w_sc =
+            data_dir + "init_param_weights_w_sc" + path_sufix;
+        init_param_path_b = data_dir + "init_param_bias_b" + path_sufix;
+        init_param_path_b_sc = data_dir + "init_param_bias_b_sc" + path_sufix;
+        opt_param_path_w = data_dir + "opt_param_weights_w" + path_sufix;
+        opt_param_path_w_sc = data_dir + "opt_param_weights_w_sc" + path_sufix;
+        opt_param_path_b = data_dir + "opt_param_bias_b" + path_sufix;
+        opt_param_path_b_sc = data_dir + "opt_param_bias_b_sc" + path_sufix;
+        forward_states_path = data_dir + "forward_hidden_states" + path_sufix;
+        backward_states_path = data_dir + "backward_hidden_states" + path_sufix;
+        input_derivative_path = data_dir + "input_derivative" + path_sufix;
+    }
+
+    std::string init_param_path_w;
+    std::string init_param_path_w_sc;
+    std::string init_param_path_b;
+    std::string init_param_path_b_sc;
+    std::string opt_param_path_w;
+    std::string opt_param_path_w_sc;
+    std::string opt_param_path_b;
+    std::string opt_param_path_b_sc;
+    std::string forward_states_path;
+    std::string backward_states_path;
+    std::string input_derivative_path;
+};
+
+/**
+ * @brief Class to store the parameters and states of the network
+ */
+class TestParamAndStates {
+   public:
+    /**
+     * @brief Construct a new Test Param And States object
+     *
+     * @param net network
+     */
+    template <typename Net>
+    TestParamAndStates(Net &net) {
+        add_weights_and_bias(weights, weights_sc, bias, bias_sc, net);
+    }
+
+    std::vector<std::vector<float> *> weights;
+    std::vector<std::vector<float> *> weights_sc;
+    std::vector<std::vector<float> *> bias;
+    std::vector<std::vector<float> *> bias_sc;
+    std::vector<std::vector<float> *> forward_states;
+    std::vector<std::vector<float> *> backward_states;
+    std::vector<std::vector<float> *> input_derivatives;
+    std::string backward_states_header;
+
+    /**
+     * @brief Write the parameters to a CSV file
+     * test_saving_paths paths to the data files
+     * init true if the parameters are initialized, false if they are optimized
+     */
+    void write_params(TestSavingPaths test_saving_paths, bool init) {
+        if (init) {
+            write_vector_to_csv(test_saving_paths.init_param_path_w, "mw,Sw",
+                                weights);
+            write_vector_to_csv(test_saving_paths.init_param_path_w_sc,
+                                "mw_sc,Sw_sc", weights_sc);
+            write_vector_to_csv(test_saving_paths.init_param_path_b, "mb,Sb",
+                                bias);
+            write_vector_to_csv(test_saving_paths.init_param_path_b_sc,
+                                "mb_sc,Sb_sc", bias_sc);
+        } else {
+            write_vector_to_csv(test_saving_paths.opt_param_path_w, "mw,Sw",
+                                weights);
+            write_vector_to_csv(test_saving_paths.opt_param_path_w_sc,
+                                "mw_sc,Sw_sc", weights_sc);
+            write_vector_to_csv(test_saving_paths.opt_param_path_b, "mb,Sb",
+                                bias);
+            write_vector_to_csv(test_saving_paths.opt_param_path_b_sc,
+                                "mb_sc,Sb_sc", bias_sc);
+        }
+    }
+
+    /**
+     * @brief Read the states from a CSV file
+     * test_saving_paths paths to the data files
+     * init true if the states are initialized, false if they are optimized
+     */
+    void read_params(TestSavingPaths test_saving_paths, bool init) {
+        if (init) {
+            read_vector_from_csv(test_saving_paths.init_param_path_w, weights);
+            read_vector_from_csv(test_saving_paths.init_param_path_w_sc,
+                                 weights_sc);
+            read_vector_from_csv(test_saving_paths.init_param_path_b, bias);
+            read_vector_from_csv(test_saving_paths.init_param_path_b_sc,
+                                 bias_sc);
+        } else {
+            read_vector_from_csv(test_saving_paths.opt_param_path_w, weights);
+            read_vector_from_csv(test_saving_paths.opt_param_path_w_sc,
+                                 weights_sc);
+            read_vector_from_csv(test_saving_paths.opt_param_path_b, bias);
+            read_vector_from_csv(test_saving_paths.opt_param_path_b_sc,
+                                 bias_sc);
+        }
+    }
+};