Repository: lhnguyen102/cuTAGI (language: python)

Issue:
Missing matrix consistency check for network architecture
Right now it is possible to train a neural net with an incorrectly defined architecture. For instance, in the image below, the input size of layer _i + 1_ must match the output size of layer _i_. However, the program does not throw any errors and attempts to train the neural net. I am not sure if it applies to layers other than Linear -- I only tested this on the regression files from the examples folder.

![image](https://github.com/user-attachments/assets/b21f44cb-c45c-4104-9a60-3eb5f290098c)



Discussion / hints:
Thanks for pointing it out. We'll work on adding the check for such case

Fix patch:
diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
new file mode 100644
index 000000000..6918fe025
--- /dev/null
+++ b/.github/pull_request_template.md
@@ -0,0 +1,23 @@
+## Description
+
+Provide a brief description of the changes introduced by this pull request. Explain the problem you're addressing or the feature you're implementing.
+
+## Changes Made
+
+Provide a detailed explanation of the changes you made in this PR. Include relevant code snippets or screenshots if necessary.
+
+## Related Issue(s)
+
+Cite any related GitHub issues or JIRA tickets that are addressed or resolved by this PR. Use the syntax `Resolves #123` or `Fixes #456`.
+
+## Checklist
+
+- [ ]  I have followed the project's coding conventions and style guidelines.
+- [ ]  I have updated the documentation, if applicable.
+- [ ]  I have rebased my branch on the latest upstream code to incorporate any changes.
+- [ ]  I have tested the changes locally.
+
+
+## Notes for Reviewers
+
+Provide any additional context or notes that the reviewers should be aware of while reviewing this PR and how to run the test.
\ No newline at end of file
diff --git a/docs/github_dev_guide.md b/docs/github_dev_guide.md
index 0f034f986..3c048d2f3 100644
--- a/docs/github_dev_guide.md
+++ b/docs/github_dev_guide.md
@@ -7,6 +7,41 @@
 4. **release/**: To prepare a new release, typically used to do tasks such as last touches and revisions.
 5. **docs/**: Used to write, modify or correct documentation.
 
+## Rebasing
+
+Follow these steps to rebase your branch onto the latest version of the `main` branch.
+
+1. **Fetch the latest updates from the remote repository:**
+   ```sh
+   git fetch
+   ```
+
+2. **Switch to the branch you want to rebase:**
+   ```sh
+   git checkout your_branch
+   ```
+
+3. **Rebase your branch onto the latest version of `main`:**
+   ```sh
+   git rebase origin/main
+   ```
+
+4. **Resolve any conflicts if any:**
+   - Manually resolve conflicts in files.
+   - After resolving, stage the changes:
+     ```sh
+     git add .
+     ```
+
+5. **Continue the rebase process:**
+   ```sh
+   git rebase --continue
+   ```
+
+6. **Verify the rebasing:**
+   - Once you see the message confirming a successful rebase, your branch is now rebased onto `main`.
+
+
 ## Pull Request Template
 ### Description
 
diff --git a/include/activation.h b/include/activation.h
index a65fba227..8da683975 100644
--- a/include/activation.h
+++ b/include/activation.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         activation.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      October 09, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <algorithm>
 #include <memory>
diff --git a/include/activation_cuda.cuh b/include/activation_cuda.cuh
index 5811b8bf9..eb77c7150 100644
--- a/include/activation_cuda.cuh
+++ b/include/activation_cuda.cuh
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         activation_cuda.cuh
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 04, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <cuda.h>
 #include <cuda_runtime.h>
diff --git a/include/base_layer.h b/include/base_layer.h
index 0b81a79e7..1c7fb5d64 100644
--- a/include/base_layer.h
+++ b/include/base_layer.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         base_layer.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      October 09, 2023
-// Updated:      April 26, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <fstream>
 #include <iostream>
diff --git a/include/base_layer_cuda.cuh b/include/base_layer_cuda.cuh
index 186a069f8..f2bf6a9ed 100644
--- a/include/base_layer_cuda.cuh
+++ b/include/base_layer_cuda.cuh
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         base_layer_cuda.cuh
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      November 29, 2023
-// Updated:      April 26, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <iostream>
 #include <memory>
diff --git a/include/base_output_updater.h b/include/base_output_updater.h
index 6da2ddb11..2ff435f9d 100644
--- a/include/base_output_updater.h
+++ b/include/base_output_updater.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         base_output_updater.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 27, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <math.h>
 
diff --git a/include/bindings/activation_bindings.h b/include/bindings/activation_bindings.h
index 70a6bd967..9d5d309ac 100644
--- a/include/bindings/activation_bindings.h
+++ b/include/bindings/activation_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         activation_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
diff --git a/include/bindings/base_layer_bindings.h b/include/bindings/base_layer_bindings.h
index b4ef3eb21..63056f2a6 100644
--- a/include/bindings/base_layer_bindings.h
+++ b/include/bindings/base_layer_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         base_layer_binding.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      December 30, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
diff --git a/include/bindings/base_layer_cuda_bindings.cuh b/include/bindings/base_layer_cuda_bindings.cuh
index 0ef422aa7..8ed8f8f5f 100644
--- a/include/bindings/base_layer_cuda_bindings.cuh
+++ b/include/bindings/base_layer_cuda_bindings.cuh
@@ -1,13 +1,3 @@
-////////////////////////////////////////////////////////////////////////////////
-// File:         base_layer_cuda_binding.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      December 30, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 
 #include <pybind11/numpy.h>
diff --git a/include/bindings/base_output_updater_bindings.h b/include/bindings/base_output_updater_bindings.h
index 562b562fa..98ce6ed6b 100644
--- a/include/bindings/base_output_updater_bindings.h
+++ b/include/bindings/base_output_updater_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         base_output_updater_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 31, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 namespace pybind11 {
 class module_;
diff --git a/include/bindings/conv2d_layer_bindings.h b/include/bindings/conv2d_layer_bindings.h
index 7ee423dca..1c60067f7 100644
--- a/include/bindings/conv2d_layer_bindings.h
+++ b/include/bindings/conv2d_layer_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         conv2d_layer_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 20, 2024
-// Updated:      January 20, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
diff --git a/include/bindings/convtranspose2d_layer_bindings.h b/include/bindings/convtranspose2d_layer_bindings.h
index e961cfb5c..5b2d3a4da 100644
--- a/include/bindings/convtranspose2d_layer_bindings.h
+++ b/include/bindings/convtranspose2d_layer_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         convtranspose2d_layer_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 19, 2024
-// Updated:      March 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
diff --git a/include/bindings/data_struct_bindings.h b/include/bindings/data_struct_bindings.h
index afcbadcd0..850395244 100644
--- a/include/bindings/data_struct_bindings.h
+++ b/include/bindings/data_struct_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         data_struct_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 31, 2023
-// Updated:      April 02, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 namespace pybind11 {
 class module_;
diff --git a/include/bindings/linear_layer_bindings.h b/include/bindings/linear_layer_bindings.h
index ae0f19c21..6f818108f 100644
--- a/include/bindings/linear_layer_bindings.h
+++ b/include/bindings/linear_layer_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         linear_layer_binding.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 29, 2023
-// Updated:      December 29, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
diff --git a/include/bindings/lstm_layer_bindings.h b/include/bindings/lstm_layer_bindings.h
index 9ff7df1ab..1e0352529 100644
--- a/include/bindings/lstm_layer_bindings.h
+++ b/include/bindings/lstm_layer_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         lstm_layer_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      December 30, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/pybind11.h>
 
diff --git a/include/bindings/main_bindings.h b/include/bindings/main_bindings.h
index 69e02e22d..4224767a5 100644
--- a/include/bindings/main_bindings.h
+++ b/include/bindings/main_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         main_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      August 21, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
diff --git a/include/bindings/norm_layer_bindings.h b/include/bindings/norm_layer_bindings.h
index 8bd7832b0..ed16c5865 100644
--- a/include/bindings/norm_layer_bindings.h
+++ b/include/bindings/norm_layer_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         norm_layer_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      February 08, 2024
-// Updated:      February 08, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/pybind11.h>
 
diff --git a/include/bindings/pooling_layer_bindings.h b/include/bindings/pooling_layer_bindings.h
index bd738ffd5..527182e51 100644
--- a/include/bindings/pooling_layer_bindings.h
+++ b/include/bindings/pooling_layer_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         pooling_layer_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 20, 2024
-// Updated:      January 20, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/pybind11.h>
 
diff --git a/include/bindings/sequential_bindings.h b/include/bindings/sequential_bindings.h
index 155902f7c..cbefc6c52 100644
--- a/include/bindings/sequential_bindings.h
+++ b/include/bindings/sequential_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         sequential_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      December 31, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 namespace pybind11 {
 class module_;
diff --git a/include/bindings/slinear_layer_bindings.h b/include/bindings/slinear_layer_bindings.h
index b2c016a95..4ffbc8772 100644
--- a/include/bindings/slinear_layer_bindings.h
+++ b/include/bindings/slinear_layer_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         slinear_layer_binding.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      August 21, 2024
-// Updated:      August 21, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
diff --git a/include/bindings/slstm_layer_bindings.h b/include/bindings/slstm_layer_bindings.h
index 8a49125c5..0417f0764 100644
--- a/include/bindings/slstm_layer_bindings.h
+++ b/include/bindings/slstm_layer_bindings.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         slstm_layer_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      August 21, 2024
-// Updated:      August 21, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <pybind11/pybind11.h>
 
diff --git a/include/bindings/utils_bindings.h b/include/bindings/utils_bindings.h
index 6640728a5..846c7f1d5 100644
--- a/include/bindings/utils_bindings.h
+++ b/include/bindings/utils_bindings.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         utils_bindings.h
-// Description:  API for Python bindings of C++/CUDA
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 31, 2024
-// Updated:      April 04, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
diff --git a/include/common.h b/include/common.h
index 18d33a12b..809c61402 100644
--- a/include/common.h
+++ b/include/common.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         common.h
-// Description:  Header file for common.h
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 15, 2022
-// Updated:      April 12, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <cmath>
 #include <fstream>
diff --git a/include/conv2d_layer.h b/include/conv2d_layer.h
index a373b27af..960b43d57 100644
--- a/include/conv2d_layer.h
+++ b/include/conv2d_layer.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         conv2d_layer.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 03, 2024
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 
 #pragma once
 #include "base_layer.h"
diff --git a/include/conv2d_layer_cuda.cuh b/include/conv2d_layer_cuda.cuh
index 8c4a615f8..3b3ed22be 100644
--- a/include/conv2d_layer_cuda.cuh
+++ b/include/conv2d_layer_cuda.cuh
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         linear_layer_cuda.cuh
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 04, 2024
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 
 #include "base_layer_cuda.cuh"
diff --git a/include/convtranspose2d_layer.h b/include/convtranspose2d_layer.h
index 73c4997e6..a710e56cf 100644
--- a/include/convtranspose2d_layer.h
+++ b/include/convtranspose2d_layer.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         convtranspose2d_layer.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 10, 2024
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include "conv2d_layer.h"
 
diff --git a/include/convtranspose2d_layer_cuda.cuh b/include/convtranspose2d_layer_cuda.cuh
index c306a3dfb..54f339387 100644
--- a/include/convtranspose2d_layer_cuda.cuh
+++ b/include/convtranspose2d_layer_cuda.cuh
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         convtranspose2d_layer_cuda.cuh
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 10, 2024
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include "base_layer_cuda.cuh"
 
diff --git a/include/cost.h b/include/cost.h
index 98d67e239..0c962c9b5 100644
--- a/include/cost.h
+++ b/include/cost.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         cost.h
-// Description:  Header file for cost function
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      February 19, 2022
-// Updated:      April 02, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <stdio.h>
 
diff --git a/include/cuda_error_checking.cuh b/include/cuda_error_checking.cuh
index c059b944a..1a5e282f3 100644
--- a/include/cuda_error_checking.cuh
+++ b/include/cuda_error_checking.cuh
@@ -1,6 +1,7 @@
 #pragma once
 #include <cuda_runtime.h>
 
+#include <chrono>
 #include <iostream>
 
 // Macro to check the last CUDA error
diff --git a/include/custom_logger.h b/include/custom_logger.h
index 18d389b4f..05c2c354f 100644
--- a/include/custom_logger.h
+++ b/include/custom_logger.h
@@ -18,4 +18,4 @@ class Logger {
 
 // Macro to simplify logging with automatic file and line number
 #define LOG(level, message) \
-    Logger::logMessage(level, __FILE__, __LINE__, message)
+    Logger::log_message(level, __FILE__, __LINE__, message)
diff --git a/include/data_struct.h b/include/data_struct.h
index e7ff6d427..0161e8652 100644
--- a/include/data_struct.h
+++ b/include/data_struct.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         data_struct.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 01, 2023
-// Updated:      April 02, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <iostream>
 #include <string>
diff --git a/include/data_struct_cuda.cuh b/include/data_struct_cuda.cuh
index ece4ae6f3..4a3ff4ed1 100644
--- a/include/data_struct_cuda.cuh
+++ b/include/data_struct_cuda.cuh
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         data_struct_cuda.cuh
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 10, 2023
-// Updated:      April 10, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <cuda.h>
 #include <cuda_runtime.h>
diff --git a/include/dataloader.h b/include/dataloader.h
index f57d32614..7f1c2a8a0 100644
--- a/include/dataloader.h
+++ b/include/dataloader.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         dataloader.h
-// Description:  Header file for dataloader
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      February 06, 2022
-// Updated:      April 02, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <assert.h>
 
diff --git a/include/debugger.h b/include/debugger.h
index b34492aa1..510cf2e54 100644
--- a/include/debugger.h
+++ b/include/debugger.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         debugger.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      February 14, 2024
-// Updated:      March 31, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include "base_output_updater.h"
 #include "sequential.h"
diff --git a/include/derivative_calcul.cuh b/include/derivative_calcul.cuh
index 7de48a2ec..ebea17c4a 100644
--- a/include/derivative_calcul.cuh
+++ b/include/derivative_calcul.cuh
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         derivative_calcul.cuh
-// Description:  Header file for derivative calculations of neural networks in
-//               cuda.
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      July 17, 2022
-// Updated:      July 27, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <cuda.h>
 #include <cuda_runtime.h>
diff --git a/include/derivative_calcul_cpu.h b/include/derivative_calcul_cpu.h
index afdd15a64..8193944f3 100644
--- a/include/derivative_calcul_cpu.h
+++ b/include/derivative_calcul_cpu.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         derivative_calculation_cpu.h
-// Description:  Header file for derivative calculation of neural networks
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      July 17, 2022
-// Updated:      July 19, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <algorithm>
 #include <iostream>
diff --git a/include/embedding_cpu.h b/include/embedding_cpu.h
index bce9cbd7f..cf62e24d4 100644
--- a/include/embedding_cpu.h
+++ b/include/embedding_cpu.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         embedding_cpu.h
-// Description:  Header file for embeddings layer
-//               (CPU version)
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      August 22, 2023
-// Updated:      September 10, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <algorithm>
 #include <iostream>
diff --git a/include/gated_residual_layer_cpu.h b/include/gated_residual_layer_cpu.h
index 0f4a214be..67e643bc4 100644
--- a/include/gated_residual_layer_cpu.h
+++ b/include/gated_residual_layer_cpu.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         gated_residual_layer_cpu.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      September 17, 2023
-// Updated:      September 18, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <vector>
 
diff --git a/include/indices.h b/include/indices.h
index d9889383e..633b162af 100644
--- a/include/indices.h
+++ b/include/indices.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         indices.h
-// Description:  Header file for indices for TAGI. Note that this is a header
-// file for indices.cpp where the detailed code are written.
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 23, 2022
-// Updated:      April 20, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include "common.h"
 
diff --git a/include/linear_layer.h b/include/linear_layer.h
index 3a1b5a988..9cddbcfc0 100644
--- a/include/linear_layer.h
+++ b/include/linear_layer.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         linear_layer.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      September 20, 2023
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <iostream>
 #include <memory>
diff --git a/include/linear_layer_cuda.cuh b/include/linear_layer_cuda.cuh
index 6c1457ee9..e1ec86732 100644
--- a/include/linear_layer_cuda.cuh
+++ b/include/linear_layer_cuda.cuh
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         linear_layer_cuda.cuh
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      November 28, 2023
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <cuda.h>
 #include <cuda_runtime.h>
diff --git a/include/lstm_layer.h b/include/lstm_layer.h
index bc5fc775c..f999d4383 100644
--- a/include/lstm_layer.h
+++ b/include/lstm_layer.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         lstm_layer.h
-// Description:  Header file for Long-Short Term Memory (LSTM) forward pass
-//               in TAGI
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 22, 2024
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <vector>
 
diff --git a/include/lstm_layer_cuda.cuh b/include/lstm_layer_cuda.cuh
index 890a43822..5f9f73053 100644
--- a/include/lstm_layer_cuda.cuh
+++ b/include/lstm_layer_cuda.cuh
@@ -1,14 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         lstm_layer.cuh
-// Description:  Header file for Long-Short Term Memory (LSTM) forward pass
-//               in TAGI
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 22, 2024
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include <cuda.h>
 #include <cuda_runtime.h>
diff --git a/include/norm_layer.h b/include/norm_layer.h
index 24ff2cae0..d11143e58 100644
--- a/include/norm_layer.h
+++ b/include/norm_layer.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         norm_layer.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 24, 2024
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 
 #include "base_layer.h"
diff --git a/include/norm_layer_cuda.cuh b/include/norm_layer_cuda.cuh
index c099fe78c..d14e97ac4 100644
--- a/include/norm_layer_cuda.cuh
+++ b/include/norm_layer_cuda.cuh
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         norm_layer_cuda.cuh
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 24, 2024
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 
 #include "base_layer_cuda.cuh"
diff --git a/include/output_layer_update_cpu.h b/include/output_layer_update_cpu.h
index ce4af3826..97de363fc 100644
--- a/include/output_layer_update_cpu.h
+++ b/include/output_layer_update_cpu.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         output_layer_update_cpu.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      October 22, 2023
-// Updated:      November 25, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <math.h>
 
diff --git a/include/output_updater_cuda.cuh b/include/output_updater_cuda.cuh
index abccb2780..f1cd6a54a 100644
--- a/include/output_updater_cuda.cuh
+++ b/include/output_updater_cuda.cuh
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         output_updater_cuda.cuh
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 27, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <cuda.h>
 #include <cuda_runtime.h>
diff --git a/include/param_init.h b/include/param_init.h
index 058f8b583..0dc75c630 100644
--- a/include/param_init.h
+++ b/include/param_init.h
@@ -1,13 +1,4 @@
 
-///////////////////////////////////////////////////////////////////////////////
-// File:         param_init.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 15, 2023
-// Updated:      March 23, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <iomanip>
 #include <iostream>
diff --git a/include/pooling_layer.h b/include/pooling_layer.h
index 276a57b6e..fae4c513c 100644
--- a/include/pooling_layer.h
+++ b/include/pooling_layer.h
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         pooling_layer.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 08, 2024
-// Updated:      March 11, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include "base_layer.h"
 
diff --git a/include/pooling_layer_cuda.cuh b/include/pooling_layer_cuda.cuh
index ece9fdc61..780d7b37a 100644
--- a/include/pooling_layer_cuda.cuh
+++ b/include/pooling_layer_cuda.cuh
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         pooling_layer_cuda.cuh
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 08, 2024
-// Updated:      January 23, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #pragma once
 #include "base_layer_cuda.cuh"
 
diff --git a/include/self_attention_cpu.h b/include/self_attention_cpu.h
index 81bb087f2..1a65585fa 100644
--- a/include/self_attention_cpu.h
+++ b/include/self_attention_cpu.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         self_attention_cpu.h
-// Description:  Header of CPU version for self attention
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 13, 2023
-// Updated:      June 09, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <cmath>
 #include <iostream>
diff --git a/include/sequential.h b/include/sequential.h
index 33ed41f39..3ca40b966 100644
--- a/include/sequential.h
+++ b/include/sequential.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         sequential.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      October 09, 2023
-// Updated:      March 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 
 #include <algorithm>
diff --git a/include/slinear_layer.h b/include/slinear_layer.h
index bb14c4283..dcb777a3c 100644
--- a/include/slinear_layer.h
+++ b/include/slinear_layer.h
@@ -1,16 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         slinear_layer.h
-// Description:  Header file for Linear layer which has
-// smoother function. It is used only as the last layer of stacked LSTM networks
-// (SLSTM).
-// Authors:     Van -Dai Vuong, Luong-Ha Nguyen & James-A. Goulet
-// Created:     August 21, 2024
-// Updated:     August 21, 2024
-// Contact:     van-dai.vuong@polymtl.ca, luongha.nguyen@gmail.com &
-// james.goulet@polymtl.ca
-// License:      This code is released under the MIT
-// License.
-////////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <iostream>
 #include <memory>
diff --git a/include/slstm_layer.h b/include/slstm_layer.h
index b82030177..60ac06537 100644
--- a/include/slstm_layer.h
+++ b/include/slstm_layer.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         slstm_layer.cpp
-// Description:  Header for Long-Short Term Memory (LSTM) which has smoother
-// function in TAGI Authors: Van-Dai Vuong, Luong-Ha Nguyen & James-A. Goulet
-// Created: August 21, 2024
-// Updated : August 21, 2024
-// Contact : van-dai.vuong@polymtl.ca, luongha.nguyen@gmail.com &
-// james.goulet@polymtl.ca
-// License:  This code is released under the MIT
 #pragma once
 #include <vector>
 
diff --git a/include/utils.h b/include/utils.h
index c769c784c..6fbe582ca 100644
--- a/include/utils.h
+++ b/include/utils.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         utils.h
-// Description:  Header file for utils
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      April 10, 2022
-// Updated:      July 19, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <string>
 #include <vector>
diff --git a/main.cpp b/main.cpp
index 8322cccca..3363b942a 100644
--- a/main.cpp
+++ b/main.cpp
@@ -7,6 +7,7 @@
 #include <iostream>
 #include <string>
 
+#include "include/custom_logger.h"
 #include "test/autoencoder/test_autoencoder_v2.h"
 #include "test/fnn/test_fnn_cpu_v2.h"
 #include "test/fnn/test_fnn_mnist_cpu.h"
@@ -22,7 +23,7 @@ int main(int argc, char* argv[]) {
     std::string user_input_file;
     std::vector<std::string> user_input_options;
     if (argc == 0) {
-        throw std::invalid_argument(
+        LOG(LogLevel::ERROR,
             "User need to provide user input file -> see README");
     } else {
         user_input_file = argv[1];
diff --git a/src/activation.cpp b/src/activation.cpp
index 642ac4f54..569f6f0c1 100644
--- a/src/activation.cpp
+++ b/src/activation.cpp
@@ -1,13 +1,5 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         activation_layer_cpu.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      October 09, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/activation.h"
+
 #ifdef USE_CUDA
 #include "activation_cuda.cuh"
 #endif
@@ -558,13 +550,6 @@ void ReLU::forward(BaseHiddenStates &input_states,
 /*
  */
 {
-    // Validate input. TODO: to be removed
-    if (input_states.size == 0) {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << ". Reason: Invalid input state size (size is 0).\n";
-        throw std::invalid_argument("Error: Invalid input state size");
-    }
-
     int start_chunk = 0;
     int end_chunk = input_states.actual_size * input_states.block_size;
     if (this->num_threads > 1) {
@@ -626,13 +611,6 @@ void Sigmoid::forward(BaseHiddenStates &input_states,
 /*
  */
 {
-    // Validate input. TODO: to be removed
-    if (input_states.size == 0) {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Input state size is zero.");
-    }
-
     // TODO: replace this function by the multiprocessing one
     int start_chunk = 0;
     int end_chunk = input_states.actual_size * input_states.block_size;
@@ -689,13 +667,6 @@ void Tanh::forward(BaseHiddenStates &input_states,
 /*
  */
 {
-    // Validate input. TODO: to be removed
-    if (input_states.size == 0) {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Input state size is zero.");
-    }
-
     // TODO: replace this function by the multiprocessing one
     int start_chunk = 0;
     int end_chunk = input_states.actual_size * input_states.block_size;
@@ -752,13 +723,6 @@ void MixtureReLU::forward(BaseHiddenStates &input_states,
 /*
  */
 {
-    // Validate input. TODO: to be removed
-    if (input_states.size == 0) {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Input state size is zero.");
-    }
-
     // TODO: replace this function by the multiprocessing one
     int start_chunk = 0;
     int end_chunk = input_states.actual_size * input_states.block_size;
@@ -816,13 +780,6 @@ void MixtureSigmoid::forward(BaseHiddenStates &input_states,
 /*
  */
 {
-    // Validate input. TODO: to be removed
-    if (input_states.size == 0) {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Input state size is zero.");
-    }
-
     // TODO: replace this function by the multiprocessing one
     int start_chunk = 0;
     int end_chunk = input_states.actual_size * input_states.block_size;
@@ -880,13 +837,6 @@ void MixtureTanh::forward(BaseHiddenStates &input_states,
 /*
  */
 {
-    // Validate input. TODO: to be removed
-    if (input_states.size == 0) {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Input state size is zero.");
-    }
-
     // TODO: replace this function by the multiprocessing one
     int start_chunk = 0;
     int end_chunk = input_states.actual_size * input_states.block_size;
@@ -943,13 +893,6 @@ void Softplus::forward(BaseHiddenStates &input_states,
 /*
  */
 {
-    // Validate input. TODO: to be removed
-    if (input_states.size == 0) {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Input state size is zero.");
-    }
-
     // TODO: replace this function by the multiprocessing one
     int start_chunk = 0;
     int end_chunk = input_states.actual_size * input_states.block_size;
@@ -1007,13 +950,6 @@ void LeakyReLU::forward(BaseHiddenStates &input_states,
 /*
  */
 {
-    // Validate input. TODO: to be removed
-    if (input_states.size == 0) {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Input state size is zero.");
-    }
-
     // TODO: replace this function by the multiprocessing one
     int start_chunk = 0;
     int end_chunk = input_states.actual_size * input_states.block_size;
@@ -1070,13 +1006,6 @@ void Softmax::forward(BaseHiddenStates &input_states,
 /*
  */
 {
-    // Validate input. TODO: to be removed
-    if (input_states.size == 0) {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Input state size is zero.");
-    }
-
     // TODO: replace this function by the multiprocessing one
     int batch_size = input_states.size / input_states.block_size;
     softmax_mean_var(input_states.mu_a, input_states.var_a,
@@ -1232,13 +1161,6 @@ void EvenExp::forward(BaseHiddenStates &input_states,
 /*
  */
 {
-    // Validate input. TODO: to be removed
-    if (input_states.size == 0) {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << ". Reason: Invalid input state size (size is 0).\n";
-        throw std::invalid_argument("Error: Invalid input state size");
-    }
-
     int start_chunk = 0;
     int end_chunk = input_states.actual_size * input_states.block_size;
     if (this->num_threads > 1) {
diff --git a/src/base_layer.cpp b/src/base_layer.cpp
index 413b34824..f851fb718 100644
--- a/src/base_layer.cpp
+++ b/src/base_layer.cpp
@@ -1,17 +1,9 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         base_layer.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      October 11, 2023
-// Updated:      April 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/base_layer.h"
 
 #include <cmath>
 
+#include "../include/custom_logger.h"
+
 InitArgs::InitArgs(size_t width, size_t height, size_t depth, int batch_size)
     : width(width), height(height), depth(depth), batch_size(batch_size) {}
 
@@ -74,10 +66,7 @@ void BaseLayer::allocate_bwd_vector(int new_size)
  */
 {
     if (new_size <= 0) {
-        throw std::invalid_argument(
-            "Error in file: " + std::string(__FILE__) +
-            " at line: " + std::to_string(__LINE__) +
-            " - Invalid size: " + std::to_string(new_size));
+        LOG(LogLevel::ERROR, "Invalid size: " + std::to_string(new_size));
     }
     this->bwd_states->set_size(new_size);
 }
@@ -139,9 +128,9 @@ void BaseLayer::update_weights()
         this->mu_w[i] +=
             delta_mu_sign * std::min(std::abs(delta_mu_w[i]), delta_bar);
         this->var_w[i] +=
-           delta_var_sign * std::min(std::abs(delta_var_w[i]), delta_bar);
+            delta_var_sign * std::min(std::abs(delta_var_w[i]), delta_bar);
         if (var_w[i] <= 0.0f) {
-            var_w[i] = 1E-5f; //TODO: replace by a parameter
+            var_w[i] = 1E-5f;  // TODO: replace by a parameter
         }
     }
 }
@@ -162,11 +151,11 @@ void BaseLayer::update_biases()
             this->mu_b[i] += delta_mu_sign *
                              std::min(std::abs(this->delta_mu_b[i]), delta_bar);
             this->var_b[i] +=
-                            delta_var_sign *
-                            std::min(std::abs(this->delta_var_b[i]), delta_bar);
+                delta_var_sign *
+                std::min(std::abs(this->delta_var_b[i]), delta_bar);
             if (var_b[i] <= 0.0f) {
-            var_b[i] = 1E-5f; //TODO: replace by a parameter
-        }
+                var_b[i] = 1E-5f;  // TODO: replace by a parameter
+            }
         }
     }
 }
@@ -188,7 +177,7 @@ factor is based on what we have tested in practice.
     if (batch_size == 1) {
         this->cap_factor_update = 0.1f;
     }
-    if (batch_size >1 && batch_size < 256) {
+    if (batch_size > 1 && batch_size < 256) {
         this->cap_factor_update = 1.0f;
     }
     if (batch_size >= 256) {
@@ -242,9 +231,7 @@ void BaseLayer::save(std::ofstream &file)
  */
 {
     if (!file.is_open()) {
-        throw std::runtime_error("Error in file: " + std::string(__FILE__) +
-                                 " at line: " + std::to_string(__LINE__) +
-                                 ". Failed to open file for saving");
+        LOG(LogLevel::ERROR, "Failed to open file for saving");
     }
 
     // Save the name length and name
@@ -273,9 +260,7 @@ void BaseLayer::load(std::ifstream &file)
  */
 {
     if (!file.is_open()) {
-        throw std::runtime_error("Error in file: " + std::string(__FILE__) +
-                                 " at line: " + std::to_string(__LINE__) +
-                                 ". Failed to open file for loading");
+        LOG(LogLevel::ERROR, "Failed to open file for loading");
     }
     // Load the name length and name
     auto layer_name = this->get_layer_info();
@@ -287,9 +272,7 @@ void BaseLayer::load(std::ifstream &file)
 
     // Check layer name
     if (layer_name != loaded_name) {
-        throw std::runtime_error("Error in file: " + std::string(__FILE__) +
-                                 " at line: " + std::to_string(__LINE__) +
-                                 ". Layer name are not match. Expected: " +
+        LOG(LogLevel::ERROR, "Layer name are not match. Expected: " +
                                  layer_name + ", Found: " + loaded_name);
     }
 
diff --git a/src/base_layer_cuda.cu b/src/base_layer_cuda.cu
index 9aaabd1fc..9f4119556 100644
--- a/src/base_layer_cuda.cu
+++ b/src/base_layer_cuda.cu
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         base_layer_cuda.cuh
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 13, 2023
-// Updated:      April 08, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/base_layer_cuda.cuh"
 #include "../include/config.h"
 #include "../include/cuda_error_checking.cuh"
diff --git a/src/base_output_updater.cpp b/src/base_output_updater.cpp
index 6bc04f1ba..7e921ab77 100644
--- a/src/base_output_updater.cpp
+++ b/src/base_output_updater.cpp
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         base_output_updater.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 27, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/base_output_updater.h"
 
 #ifdef USE_CUDA
diff --git a/src/bindings/activation_bindings.cpp b/src/bindings/activation_bindings.cpp
index dbe34953c..b4911518c 100644
--- a/src/bindings/activation_bindings.cpp
+++ b/src/bindings/activation_bindings.cpp
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         activation_bindings.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 
 #include "../include/bindings/activation_bindings.h"
 
diff --git a/src/bindings/base_layer_cuda_bindings.cu b/src/bindings/base_layer_cuda_bindings.cu
index ed7b4a9c4..a773ab9ad 100644
--- a/src/bindings/base_layer_cuda_bindings.cu
+++ b/src/bindings/base_layer_cuda_bindings.cu
@@ -1,13 +1,3 @@
-////////////////////////////////////////////////////////////////////////////////
-// File:         base_layer_cuda_binding.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      December 30, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/linear_layer_cuda_bindings.h"
 
 void base_layer_cuda_binding(pybind11::module_& modo)
diff --git a/src/bindings/base_output_updater_bindings.cpp b/src/bindings/base_output_updater_bindings.cpp
index 7640c540a..4320f95fb 100644
--- a/src/bindings/base_output_updater_bindings.cpp
+++ b/src/bindings/base_output_updater_bindings.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         base_output_updater_bindings.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 31, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/base_output_updater_bindings.h"
 
 #include <pybind11/numpy.h>
diff --git a/src/bindings/conv2d_layer_bindings.cpp b/src/bindings/conv2d_layer_bindings.cpp
index c5497d100..f69446dac 100644
--- a/src/bindings/conv2d_layer_bindings.cpp
+++ b/src/bindings/conv2d_layer_bindings.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         conv2d_layer_bindings.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 20, 2024
-// Updated:      January 20, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/conv2d_layer_bindings.h"
 
 #include "../include/conv2d_layer.h"
diff --git a/src/bindings/convtranspose2d_layer_bindings.cpp b/src/bindings/convtranspose2d_layer_bindings.cpp
index 24c6c13ad..f0d2b2a90 100644
--- a/src/bindings/convtranspose2d_layer_bindings.cpp
+++ b/src/bindings/convtranspose2d_layer_bindings.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         convtranspose2d_layer_bindings.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 19, 2024
-// Updated:      March 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/convtranspose2d_layer_bindings.h"
 
 #include "../include/base_layer.h"
diff --git a/src/bindings/data_struct_bindings.cpp b/src/bindings/data_struct_bindings.cpp
index 163e33595..b0df89a5a 100644
--- a/src/bindings/data_struct_bindings.cpp
+++ b/src/bindings/data_struct_bindings.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         data_struct_bindings.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 31, 2023
-// Updated:      April 11, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/data_struct_bindings.h"
 
 #include <pybind11/numpy.h>
diff --git a/src/bindings/linear_layer_bindings.cpp b/src/bindings/linear_layer_bindings.cpp
index 92cda1a6f..3bdbe1dbe 100644
--- a/src/bindings/linear_layer_bindings.cpp
+++ b/src/bindings/linear_layer_bindings.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         linear_layer_binding.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      March 09, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/linear_layer_bindings.h"
 
 void bind_linear_layer(pybind11::module_& modo)
diff --git a/src/bindings/lstm_layer_bindings.cpp b/src/bindings/lstm_layer_bindings.cpp
index 0f6393d33..ecb754f1e 100644
--- a/src/bindings/lstm_layer_bindings.cpp
+++ b/src/bindings/lstm_layer_bindings.cpp
@@ -1,14 +1,3 @@
-
-///////////////////////////////////////////////////////////////////////////////
-// File:         lstm_layer_binding.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      December 30, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/lstm_layer_bindings.h"
 
 #include <pybind11/numpy.h>
diff --git a/src/bindings/main_bindings.cpp b/src/bindings/main_bindings.cpp
index 5838b6c08..9ae8507f3 100644
--- a/src/bindings/main_bindings.cpp
+++ b/src/bindings/main_bindings.cpp
@@ -1,14 +1,3 @@
-
-///////////////////////////////////////////////////////////////////////////////
-// File:         main_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/main_bindings.h"
 
 PYBIND11_MODULE(cutagi, modo) {
diff --git a/src/bindings/norm_layer_bindings.cpp b/src/bindings/norm_layer_bindings.cpp
index 3702c0bc8..db21a2017 100644
--- a/src/bindings/norm_layer_bindings.cpp
+++ b/src/bindings/norm_layer_bindings.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         norm_layer_bindings.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      February 08, 2024
-// Updated:      April 12, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/norm_layer_bindings.h"
 
 #include <pybind11/numpy.h>
diff --git a/src/bindings/pooling_layer_bindings.cpp b/src/bindings/pooling_layer_bindings.cpp
index c3f72356d..c8586f64f 100644
--- a/src/bindings/pooling_layer_bindings.cpp
+++ b/src/bindings/pooling_layer_bindings.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         pooling_layer_bindings.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 20, 2024
-// Updated:      January 20, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/pooling_layer_bindings.h"
 
 #include <pybind11/numpy.h>
diff --git a/src/bindings/sequential_bindings.cpp b/src/bindings/sequential_bindings.cpp
index a16b904aa..7a05322ac 100644
--- a/src/bindings/sequential_bindings.cpp
+++ b/src/bindings/sequential_bindings.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         sequential_bindings.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 30, 2023
-// Updated:      March 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/sequential_bindings.h"
 
 #include <pybind11/numpy.h>
diff --git a/src/bindings/slinear_layer_bindings.cpp b/src/bindings/slinear_layer_bindings.cpp
index e4ec44afd..e9102c98d 100644
--- a/src/bindings/slinear_layer_bindings.cpp
+++ b/src/bindings/slinear_layer_bindings.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         slinear_layer_binding.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      August 21, 2024
-// Updated:      August 21, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/slinear_layer_bindings.h"
 
 void bind_slinear_layer(pybind11::module_& modo)
diff --git a/src/bindings/slstm_layer_bindings.cpp b/src/bindings/slstm_layer_bindings.cpp
index 038a56e0c..6de1ffb18 100644
--- a/src/bindings/slstm_layer_bindings.cpp
+++ b/src/bindings/slstm_layer_bindings.cpp
@@ -1,14 +1,3 @@
-
-///////////////////////////////////////////////////////////////////////////////
-// File:         slstm_layer_binding.h
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      August 21, 2024
-// Updated:      August 21, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/slstm_layer_bindings.h"
 
 #include <pybind11/numpy.h>
diff --git a/src/bindings/utils_bindings.cpp b/src/bindings/utils_bindings.cpp
index dfc23fda3..121dd7931 100644
--- a/src/bindings/utils_bindings.cpp
+++ b/src/bindings/utils_bindings.cpp
@@ -1,14 +1,3 @@
-
-///////////////////////////////////////////////////////////////////////////////
-// File:         utils_bindings.cpp
-// Description:  API for Python bindings of C++/CUDA
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 31, 2024
-// Updated:      April 04, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
-
 #include "../include/bindings/utils_bindings.h"
 
 Utils::Utils() {};
diff --git a/src/common.cpp b/src/common.cpp
index cf6ba4503..e970cfdd7 100644
--- a/src/common.cpp
+++ b/src/common.cpp
@@ -1,14 +1,7 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         common.cpp
-// Description:  Common function used for computing indices for TAGI
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 15, 2022
-// Updated:      April 12, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
 #include "../include/common.h"
 
+#include "../include/custom_logger.h"
+
 std::string get_current_dir() {
     char buff[FILENAME_MAX];  // create string buffer to hold path
     GetCurrentDir(buff, FILENAME_MAX);
@@ -325,7 +318,7 @@ float normpdf_cpu(float x, float mu, float sigma)
 /*Probability density function of Normal distribution*/
 {
     if (sigma < 0.0f) {
-        throw std::invalid_argument("Sigma value is negative");
+        LOG(LogLevel::ERROR, "Sigma value is negative");
     }
     const float PI = 3.14159265358979323846f;
     float prob_pdf = (1 / (sigma * pow(2 * PI, 0.5))) *
diff --git a/src/conv2d_layer.cpp b/src/conv2d_layer.cpp
index 7fb1d59d2..6f37d288f 100644
--- a/src/conv2d_layer.cpp
+++ b/src/conv2d_layer.cpp
@@ -1,15 +1,6 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         conv2d_layer.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 04, 2024
-// Updated:      August 06, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/conv2d_layer.h"
 
+#include "../include/custom_logger.h"
 #include "../include/indices.h"
 #include "../include/param_init.h"
 #ifdef USE_CUDA
@@ -128,6 +119,14 @@ void Conv2d::forward(BaseHiddenStates &input_states,
 /*
  */
 {
+    // Checkout input size
+    if (this->input_size != input_states.actual_size) {
+        std::string message =
+            "Input size mismatch: " + std::to_string(this->input_size) +
+            " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     int batch_size = input_states.block_size;
     this->set_cap_factor_udapte(batch_size);
 
@@ -615,10 +614,7 @@ std::tuple<int, int> compute_downsample_img_size_v2(int kernel, int stride,
         wo = nom_w / stride + 1;
         ho = nom_h / stride + 1;
     } else {
-        throw std::invalid_argument(
-            "Error in file: " + std::string(__FILE__) +
-            " at line: " + std::to_string(__LINE__) +
-            ". Invalid hyperparameters for conv2d layer");
+        LOG(LogLevel::ERROR, "Invalid hyperparameters for conv2d layer");
     }
 
     return {wo, ho};
diff --git a/src/conv2d_layer_cuda.cu b/src/conv2d_layer_cuda.cu
index 58e1ee5c7..ab7ddcc86 100644
--- a/src/conv2d_layer_cuda.cu
+++ b/src/conv2d_layer_cuda.cu
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         conv2d_layer_cuda.cu
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 04, 2024
-// Updated:      July 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 
 #include <cuda.h>
 #include <cuda_runtime.h>
@@ -23,6 +14,7 @@
 #include "../include/conv2d_layer.h"
 #include "../include/conv2d_layer_cuda.cuh"
 #include "../include/cuda_error_checking.cuh"
+#include "../include/custom_logger.h"
 #include "../include/param_init.h"
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -395,6 +387,14 @@ void Conv2dCuda::forward(BaseHiddenStates &input_states,
 /*
  */
 {
+    // Checkout input size
+    if (this->input_size != input_states.actual_size) {
+        std::string message =
+            "Input size mismatch: " + std::to_string(this->input_size) +
+            " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     // New poitner will point to the same memory location when casting
     HiddenStateCuda *cu_input_states =
         dynamic_cast<HiddenStateCuda *>(&input_states);
diff --git a/src/convtranspose2d_layer.cpp b/src/convtranspose2d_layer.cpp
index e634e947e..ac78d78ed 100644
--- a/src/convtranspose2d_layer.cpp
+++ b/src/convtranspose2d_layer.cpp
@@ -1,25 +1,15 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         convtranspose2d_layer.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 10, 2024
-// Updated:      August 06, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/convtranspose2d_layer.h"
-#ifdef USE_CUDA
-#include <cuda.h>
-#include <cuda_runtime.h>
-#endif
 
 #include "../include/base_layer.h"
 #include "../include/common.h"
+#include "../include/custom_logger.h"
 #include "../include/indices.h"
 #include "../include/param_init.h"
 
 #ifdef USE_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+
 #include "../include/convtranspose2d_layer_cuda.cuh"
 #endif
 #include <thread>
@@ -447,10 +437,7 @@ std::tuple<int, int> compute_upsample_img_size_v2(int kernel, int stride,
     }
 
     if (nom_w % stride != 0 || nom_h % stride != 0) {
-        throw std::invalid_argument(
-            "Error in file: " + std::string(__FILE__) +
-            " at line: " + std::to_string(__LINE__) +
-            ". Invalid hyperparameters for ConvTranspose2d layer");
+        LOG(LogLevel::ERROR, "Invalid hyperparameters for conv2d layer");
     }
 
     return {wo, ho};
@@ -572,6 +559,14 @@ void ConvTranspose2d::forward(BaseHiddenStates &input_states,
 /*
  */
 {
+    // Checkout input size
+    if (this->input_size != input_states.actual_size) {
+        std::string message =
+            "Input size mismatch: " + std::to_string(this->input_size) +
+            " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     int batch_size = input_states.block_size;
     this->set_cap_factor_udapte(batch_size);
 
diff --git a/src/convtranspose2d_layer_cuda.cu b/src/convtranspose2d_layer_cuda.cu
index e836f2f97..eb99a3294 100644
--- a/src/convtranspose2d_layer_cuda.cu
+++ b/src/convtranspose2d_layer_cuda.cu
@@ -1,12 +1,9 @@
-///////////////////////////////////////////////////////////////////////////////
-// This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/common.h"
 #include "../include/config.h"
 #include "../include/convtranspose2d_cuda_kernel.cuh"
 #include "../include/convtranspose2d_layer.h"
 #include "../include/convtranspose2d_layer_cuda.cuh"
+#include "../include/custom_logger.h"
 #include "../include/param_init.h"
 
 ConvTranspose2dCuda::ConvTranspose2dCuda(
@@ -156,9 +153,7 @@ void ConvTranspose2dCuda::allocate_convtranspose_index()
 
     cudaError_t error = cudaGetLastError();
     if (error != cudaSuccess) {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Device memory allocation.");
+        LOG(LogLevel::ERROR, "Device memory allocation.");
     }
 }
 
@@ -184,10 +179,8 @@ void ConvTranspose2dCuda::convtranspose_index_to_device()
 
     cudaError_t error = cudaGetLastError();
     if (error != cudaSuccess) {
-        fprintf(stderr, "CUDA Error: %s\n", cudaGetErrorString(error));
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Host to device.");
+        LOG(LogLevel::ERROR,
+            std::string("Host to device. ") + cudaGetErrorString(error));
     }
 }
 
@@ -197,6 +190,14 @@ void ConvTranspose2dCuda::forward(BaseHiddenStates &input_states,
 /*
  */
 {
+    // Checkout input size
+    if (this->input_size != input_states.actual_size) {
+        std::string message =
+            "Input size mismatch: " + std::to_string(this->input_size) +
+            " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     // New poitner will point to the same memory location when casting
     HiddenStateCuda *cu_input_states =
         dynamic_cast<HiddenStateCuda *>(&input_states);
diff --git a/src/cost.cpp b/src/cost.cpp
index 2e885ecb9..26d0050d6 100644
--- a/src/cost.cpp
+++ b/src/cost.cpp
@@ -1,14 +1,7 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         cost.cpp
-// Description:  Cost functions for TAGI
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 01, 2022
-// Updated:      April 02, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/cost.h"
 
+#include "../include/custom_logger.h"
+
 void fliplr(std::vector<int> &v)
 /* Flip an array from left to right
  *
@@ -337,9 +330,8 @@ float mean_squared_error(std::vector<float> &pred, std::vector<float> &obs)
 */
 {
     if (pred.size() != obs.size()) {
-        throw std::invalid_argument(
-            "Prediciton and observation does not have the same lenght - "
-            "cost.cpp");
+        LOG(LogLevel::ERROR,
+            "Prediciton and observation does not have the same.");
     }
     float sum = 0;
     for (int i = 0; i < pred.size(); i++) {
@@ -365,8 +357,7 @@ float avg_univar_log_lik(std::vector<float> &x, std::vector<float> &mu,
 */
 {
     if (x.size() == 0 || mu.size() == 0 || sigma.size() == 0) {
-        throw std::invalid_argument(
-            "Invalid inputs for normal density - cost.cpp");
+        LOG(LogLevel::ERROR, "Invalid inputs for normal density");
     }
     float sum = 0;
     float PI_C = 3.141592653f;
diff --git a/src/data_struct.cpp b/src/data_struct.cpp
index da43068e3..41f896c4a 100644
--- a/src/data_struct.cpp
+++ b/src/data_struct.cpp
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         data_struct.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 10, 2023
-// Updated:      April 11, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/data_struct.h"
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/src/data_struct_cuda.cu b/src/data_struct_cuda.cu
index f7d697ace..d5d8582ee 100644
--- a/src/data_struct_cuda.cu
+++ b/src/data_struct_cuda.cu
@@ -1,15 +1,6 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         data_struct_cuda.cu
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 10, 2023
-// Updated:      July 12, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/config.h"
 #include "../include/cuda_error_checking.cuh"
+#include "../include/custom_logger.h"
 #include "../include/data_struct_cuda.cuh"
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -158,9 +149,7 @@ void HiddenStateCuda::swap(BaseHiddenStates &other) {
         std::swap(this->d_var_a, cu_other->d_var_a);
         std::swap(this->d_jcb, cu_other->d_jcb);
     } else {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Swap input invalid.");
+        LOG(LogLevel::ERROR, "Swap input invalid.");
     }
 }
 
@@ -176,9 +165,7 @@ void HiddenStateCuda::copy_from(const BaseHiddenStates &source, int num_data)
         dynamic_cast<const HiddenStateCuda *>(&source);
 
     if (!cu_source) {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Invalid source.");
+        LOG(LogLevel::ERROR, "Invalid source.");
     }
 
     cudaMemcpy(this->d_mu_a, cu_source->d_mu_a, num_data * sizeof(float),
@@ -279,9 +266,7 @@ void DeltaStateCuda::copy_from(const BaseDeltaStates &source, int num_data)
         dynamic_cast<const DeltaStateCuda *>(&source);
 
     if (!cu_source) {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Invalid source.");
+        LOG(LogLevel::ERROR, "Invalid source.");
     }
 
     cudaMemcpy(this->d_delta_mu, cu_source->d_delta_mu,
@@ -317,9 +302,7 @@ void DeltaStateCuda::swap(BaseDeltaStates &other) {
         std::swap(this->d_delta_mu, cu_other->d_delta_mu);
         std::swap(this->d_delta_var, cu_other->d_delta_var);
     } else {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Swap input invalid.");
+        LOG(LogLevel::ERROR, "Swap input invalid.");
     }
 }
 
diff --git a/src/dataloader.cpp b/src/dataloader.cpp
index bd8f2ba7f..8c2a1482e 100644
--- a/src/dataloader.cpp
+++ b/src/dataloader.cpp
@@ -1,15 +1,7 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         dataloader.cpp
-// Description:  Load different batches of data to network
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      February 06, 2022
-// Updated:      April 04, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
-
 #include "../include/dataloader.h"
 
+#include "../include/custom_logger.h"
+
 std::vector<int> create_range(int N)
 /*
  * Create a range vector.
@@ -105,9 +97,7 @@ std::vector<float> load_mnist_images(std::string image_file, int num)
     std::ifstream data_file(image_file.c_str(),
                             std::ios::in | std::ios::binary);
     if (data_file.fail()) {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Image files do not exist.");
+        LOG(LogLevel::ERROR, "Image files do not exist.");
     }
 
     // Check the magic number (see http://yann.lecun.com/exdb/mnist/)
@@ -169,9 +159,7 @@ std::vector<int> load_mnist_labels(std::string label_file, int num)
                             std::ios::in | std::ios::binary);
 
     if (data_file.fail()) {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Label files do not exist.");
+        LOG(LogLevel::ERROR, "Label files do not exist.");
     }
 
     // Check the magic number
@@ -229,9 +217,7 @@ std::tuple<std::vector<float>, std::vector<int>> load_cifar_images(
     std::ifstream data_file(image_file.c_str(),
                             std::ios::in | std::ios::binary);
     if (data_file.fail()) {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Image files do not exist.");
+        LOG(LogLevel::ERROR, "Image files do not exist.");
     }
 
     int n_size = 10000;
@@ -318,9 +304,7 @@ ImageData get_images_v2(std::string data_name,
         image_data.image_len = 32 * 32 * 3;
 
     } else {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Dataset does not exist.");
+        LOG(LogLevel::ERROR, "Dataset does not exist.");
     }
 
     // Convert label to hierarchical softmax
@@ -603,7 +587,7 @@ void create_rolling_windows(std::vector<float> &data,
         (data.size() / num_features - num_input_ts - num_output_ts) / stride +
         1;
     if (num_samples < 0) {
-        throw std::invalid_argument("Could not window time series data");
+        LOG(LogLevel::ERROR, "Could not window time series data");
     }
     int num_outputs = output_col.size();
 
diff --git a/src/debugger.cpp b/src/debugger.cpp
index 28362079d..3b445f074 100644
--- a/src/debugger.cpp
+++ b/src/debugger.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         debugger.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      February 14, 2024
-// Updated:      March 18, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/debugger.h"
 
 #include <memory>
diff --git a/src/derivative_calcul.cu b/src/derivative_calcul.cu
index 541110843..eafa93dfa 100644
--- a/src/derivative_calcul.cu
+++ b/src/derivative_calcul.cu
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         derivative_calcul.cu
-// Description:  Calculate derivatives of neural networks
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      July 17, 2022
-// Updated:      July 30, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/derivative_calcul.cuh"
 
 __global__ void compute_node_derv_mean_var_fc(float const *mw, float const *Sw,
diff --git a/src/derivative_calcul_cpu.cpp b/src/derivative_calcul_cpu.cpp
index c1253d845..c0dbd998d 100644
--- a/src/derivative_calcul_cpu.cpp
+++ b/src/derivative_calcul_cpu.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         derivative_calcul_cpu.cpp
-// Description:  Calculate derivatives of neural networks
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      July 12, 2022
-// Updated:      July 30, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/derivative_calcul_cpu.h"
 
 void compute_node_derv_mean_var_fc_cpu(
diff --git a/src/embedding_cpu.cpp b/src/embedding_cpu.cpp
index b7db285dc..a7af760d2 100644
--- a/src/embedding_cpu.cpp
+++ b/src/embedding_cpu.cpp
@@ -1,14 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         embedding_cpu.cpp
-// Description:  embeddings layer
-//               (CPU version)
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      August 23, 2023
-// Updated:      September 06, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
-
 #include "../include/embedding_cpu.h"
 
 std::tuple<std::vector<float>, std::vector<float>> get_embedding_values(
diff --git a/src/gated_residual_layer_cpu.cpp b/src/gated_residual_layer_cpu.cpp
index 32cd33a55..f1b304cc8 100644
--- a/src/gated_residual_layer_cpu.cpp
+++ b/src/gated_residual_layer_cpu.cpp
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         gated_residual_layer_cpu.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      September 17, 2023
-// Updated:      September 17, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/gated_residual_layer_cpu.h"
 
 int get_number_of_params(GatedResidualProp &prop)
diff --git a/src/indices.cpp b/src/indices.cpp
index ad0c8a6b0..ce2313ceb 100644
--- a/src/indices.cpp
+++ b/src/indices.cpp
@@ -1,17 +1,9 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         indices.cpp
-// Description:  Indices for TAGI
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 12, 2021
-// Updated:      May 07, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
-
 #include "../include/indices.h"
 
 #include <algorithm>
 
+#include "../include/custom_logger.h"
+
 ////////////////////////////
 // IMAGE CONSTRUCTION
 ////////////////////////////
@@ -253,7 +245,7 @@ std::vector<int> get_receptive_field(std::vector<int> &img,
     // Check validity for recetive field
 
     if ((hi - kernel) / stride + 1 != wo && kernel != 1) {
-        throw std::invalid_argument("Kernel size is invalid");
+        LOG(LogLevel::ERROR, "Invalid receptive field");
     }
     // Get reference kernel
     std::vector<int> ref_kernel = get_ref_kernel(img, kernel, wi);
@@ -658,7 +650,7 @@ std::vector<int> get_FCzwa_2_idx(std::vector<int> &Fmwa_2_idx, int pad,
 
     // Reorganize indices as reference indices
     if (idx_ref.size() != tmp.size()) {
-        throw std::invalid_argument(
+        LOG(LogLevel::ERROR,
             "Size of reference indices is not equal to size of FCzwa 2");
     }
 
diff --git a/src/layer_block.cpp b/src/layer_block.cpp
index 110d107fa..87f308b5f 100644
--- a/src/layer_block.cpp
+++ b/src/layer_block.cpp
@@ -1,5 +1,7 @@
 #include "../include/layer_block.h"
 
+#include "../include/custom_logger.h"
+
 #ifdef USE_CUDA
 #include "../include/base_layer_cuda.cuh"
 #endif
@@ -37,9 +39,7 @@ it will be corrected at the first run in the forward pass.
     } else if (this->device.compare("cuda") == 0) {
         this->layers.push_back(std::move(layer->to_cuda()));
     } else {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Invalid device: [" + this->device + "]");
+        LOG(LogLevel::ERROR, "Invalid device: [" + this->device + "]");
     }
 }
 
diff --git a/src/linear_layer.cpp b/src/linear_layer.cpp
index 1011db7f8..992ba8e52 100644
--- a/src/linear_layer.cpp
+++ b/src/linear_layer.cpp
@@ -1,15 +1,7 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         linear_layer.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      September 20, 2023
-// Updated:      August 06, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/linear_layer.h"
 
 #include "../include/common.h"
+#include "../include/custom_logger.h"
 
 #ifdef USE_CUDA
 #include "../include/linear_layer_cuda.cuh"
@@ -538,6 +530,14 @@ void Linear::forward(BaseHiddenStates &input_states,
     int batch_size = input_states.block_size;
     this->set_cap_factor_udapte(batch_size);
 
+    // Checkout input size
+    if (this->input_size != input_states.actual_size) {
+        std::string message =
+            "Input size mismatch: " + std::to_string(this->input_size) +
+            " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     // Forward pass
     if (this->num_threads > 1) {
         linear_fwd_mean_var_mp(this->mu_w, this->var_w, this->mu_b, this->var_b,
diff --git a/src/linear_layer_cuda.cu b/src/linear_layer_cuda.cu
index 0df6db50a..c73b77bd6 100644
--- a/src/linear_layer_cuda.cu
+++ b/src/linear_layer_cuda.cu
@@ -1,15 +1,7 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         linear_layer_cuda.cu
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 03, 2023
-// Updated:      March 09, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include <cstdint>
 
 #include "../include/config.h"
+#include "../include/custom_logger.h"
 #include "../include/linear_layer.h"
 #include "../include/linear_layer_cuda.cuh"
 
@@ -1762,6 +1754,14 @@ void LinearCuda::forward(BaseHiddenStates &input_states,
     // TempStateCuda *cu_temp_states = dynamic_cast<TempStateCuda
     // *>(&temp_states);
 
+    // Checkout input size
+    if (this->input_size != input_states.actual_size) {
+        std::string message =
+            "Input size mismatch: " + std::to_string(this->input_size) +
+            " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     int batch_size = input_states.block_size;
 
     this->set_cap_factor_udapte(batch_size);
diff --git a/src/lstm_layer.cpp b/src/lstm_layer.cpp
index a47d0d15b..ec64bc704 100644
--- a/src/lstm_layer.cpp
+++ b/src/lstm_layer.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         lstm_layer.cpp
-// Description:  Header file for Long-Short Term Memory (LSTM) forward pass
-//               in TAGI
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 22, 2024
-// Updated:      March 27, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 
 #include "../include/lstm_layer.h"
 
@@ -16,6 +6,7 @@
 
 #include "../include/activation.h"
 #include "../include/common.h"
+#include "../include/custom_logger.h"
 #include "../include/param_init.h"
 
 #ifdef USE_CUDA
@@ -1260,6 +1251,14 @@ void LSTM::forward(BaseHiddenStates &input_states,
 /*
  */
 {
+    // Checkout input size
+    if (this->input_size * this->seq_len != input_states.actual_size) {
+        std::string message = "Input size mismatch: " +
+                              std::to_string(this->input_size * this->seq_len) +
+                              " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     int batch_size = input_states.block_size;
     this->set_cap_factor_udapte(batch_size);
 
diff --git a/src/lstm_layer_cuda.cu b/src/lstm_layer_cuda.cu
index bfe2b8f59..b1e29e1fe 100644
--- a/src/lstm_layer_cuda.cu
+++ b/src/lstm_layer_cuda.cu
@@ -1,15 +1,6 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         lstm_layer_cuda.cu
-// Description:  Header file for Long-Short Term Memory (LSTM) forward pass
-//               in TAGI
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 22, 2024
-// Updated:      March 31, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 
 #include "../include/activation_cuda.cuh"
+#include "../include/custom_logger.h"
 #include "../include/lstm_layer.h"
 #include "../include/lstm_layer_cuda.cuh"
 #include "../include/param_init.h"
@@ -814,6 +805,14 @@ void LSTMCuda::forward(BaseHiddenStates &input_states,
 /*
  */
 {
+    // Checkout input size
+    if (this->input_size * this->seq_len != input_states.actual_size) {
+        std::string message = "Input size mismatch: " +
+                              std::to_string(this->input_size * this->seq_len) +
+                              " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     // New poitner will point to the same memory location when casting
     HiddenStateCuda *cu_input_states =
         dynamic_cast<HiddenStateCuda *>(&input_states);
diff --git a/src/norm_layer.cpp b/src/norm_layer.cpp
index 851c00bd8..c8627f5ce 100644
--- a/src/norm_layer.cpp
+++ b/src/norm_layer.cpp
@@ -1,14 +1,7 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         norm_layer.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 24, 2024
-// Updated:      August 6th, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/norm_layer.h"
 
+#include "../include/custom_logger.h"
+
 #ifdef USE_CUDA
 #include "../include/norm_layer_cuda.cuh"
 #endif
@@ -1485,6 +1478,14 @@ void LayerNorm::forward(BaseHiddenStates &input_states,
                         BaseTempStates &temp_states)
 /**/
 {
+    // Checkout input size
+    if (this->input_size != input_states.actual_size) {
+        std::string message =
+            "Input size mismatch: " + std::to_string(this->input_size) +
+            " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     int batch_size = input_states.block_size;
     if (this->_batch_size != batch_size) {
         this->_batch_size = batch_size;
diff --git a/src/norm_layer_cuda.cu b/src/norm_layer_cuda.cu
index 1adc5aa89..5bae6b81e 100644
--- a/src/norm_layer_cuda.cu
+++ b/src/norm_layer_cuda.cu
@@ -1,14 +1,6 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         norm_layer_cuda.cu
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 24, 2024
-// Updated:      July 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/config.h"
 #include "../include/cuda_error_checking.cuh"
+#include "../include/custom_logger.h"
 #include "../include/norm_layer.h"
 #include "../include/norm_layer_cuda.cuh"
 
@@ -1069,6 +1061,14 @@ void LayerNormCuda::forward(BaseHiddenStates &input_states,
 /*
  */
 {
+    // Checkout input size
+    if (this->input_size != input_states.actual_size) {
+        std::string message =
+            "Input size mismatch: " + std::to_string(this->input_size) +
+            " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     // New poitner will point to the same memory location when casting
     HiddenStateCuda *cu_input_states =
         dynamic_cast<HiddenStateCuda *>(&input_states);
diff --git a/src/output_layer_update_cpu.cpp b/src/output_layer_update_cpu.cpp
index a082ea357..74c3d47f2 100644
--- a/src/output_layer_update_cpu.cpp
+++ b/src/output_layer_update_cpu.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         output_layer_update_cpu.cpu
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      October 22, 2023
-// Updated:      November 17, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/output_layer_update_cpu.h"
 
 void update_output_delta_z(BaseHiddenStates &last_layer_states,
diff --git a/src/output_updater_cuda.cu b/src/output_updater_cuda.cu
index 886db823d..5b364bce9 100644
--- a/src/output_updater_cuda.cu
+++ b/src/output_updater_cuda.cu
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         output_updater_cuda.cu
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 27, 2023
-// Updated:      August 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/output_updater_cuda.cuh"
 
 __global__ void update_delta_z_using_indices_cuda(
diff --git a/src/param_init.cpp b/src/param_init.cpp
index 85960e468..23b16ae5b 100644
--- a/src/param_init.cpp
+++ b/src/param_init.cpp
@@ -1,14 +1,7 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         param_init.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      December 15, 2023
-// Updated:      March 23, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/param_init.h"
 
+#include "../include/custom_logger.h"
+
 float he_init(float fan_in)
 
 /* He initialization for neural networks. Further details can be found in
@@ -146,10 +139,8 @@ init_weight_bias_linear(const std::string &init_method, const float gain_w,
                init_method.compare("he") == 0) {
         scale = he_init(input_size);
     } else {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Initial parameter method [" +
-                                    init_method + "] is not supported.");
+        LOG(LogLevel::ERROR,
+            "Initial parameter method [" + init_method + "] is not supported.");
     }
 
     // Initalize weights & biases
@@ -182,10 +173,8 @@ init_weight_bias_conv2d(const size_t kernel_size, const size_t in_channels,
                init_method.compare("he") == 0) {
         scale = he_init(fan_in);
     } else {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Initial parameter method [" +
-                                    init_method + "] is not supported.");
+        LOG(LogLevel::ERROR,
+            "Initial parameter method [" + init_method + "] is not supported.");
     }
 
     // Initalize weights & biases
@@ -213,10 +202,8 @@ init_weight_bias_lstm(const std::string &init_method, const float gain_w,
                init_method.compare("he") == 0) {
         scale = he_init(input_size + output_size);
     } else {
-        std::cerr << "Error in file: " << __FILE__ << " at line: " << __LINE__
-                  << std::endl;
-        throw std::invalid_argument("Error: Initial parameter method [" +
-                                    init_method + "] is not supported.");
+        LOG(LogLevel::ERROR,
+            "Initial parameter method [" + init_method + "] is not supported.");
     }
 
     // Initalize weights & biases for 4 gates
diff --git a/src/pooling_layer.cpp b/src/pooling_layer.cpp
index 0f19804a5..652832488 100644
--- a/src/pooling_layer.cpp
+++ b/src/pooling_layer.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         pooling_layer.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 08, 2024
-// Updated:      January 23, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/pooling_layer.h"
 
 #include "../include/conv2d_layer.h"
diff --git a/src/pooling_layer_cuda.cu b/src/pooling_layer_cuda.cu
index 6bc94646a..468f0ae39 100644
--- a/src/pooling_layer_cuda.cu
+++ b/src/pooling_layer_cuda.cu
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         pooling_layer_cuda.cu
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      January 08, 2024
-// Updated:      July 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include <cuda.h>
 #include <cuda_runtime.h>
 
diff --git a/src/resnet_block_cuda.cu b/src/resnet_block_cuda.cu
index 08102029c..ed496b383 100644
--- a/src/resnet_block_cuda.cu
+++ b/src/resnet_block_cuda.cu
@@ -1,5 +1,6 @@
 #include <cuda.h>
 
+#include "../include/custom_logger.h"
 #include "../include/resnet_block.h"
 #include "../include/resnet_block_cuda.cuh"
 
@@ -191,9 +192,7 @@ void ResNetBlockCuda::set_cuda_threads(int num)
         if (layer_block) {
             layer_block->set_cuda_threads(num);
         } else {
-            throw std::invalid_argument(
-                "Error in file: " + std::string(__FILE__) + " at line: " +
-                std::to_string(__LINE__) + ". Set cuda threads.");
+            LOG(LogLevel::ERROR, "Set cuda threads.");
         }
     }
 
@@ -203,9 +202,7 @@ void ResNetBlockCuda::set_cuda_threads(int num)
         if (cu_shortcut) {
             cu_shortcut->set_cuda_threads(num);
         } else {
-            throw std::invalid_argument(
-                "Error in file: " + std::string(__FILE__) + " at line: " +
-                std::to_string(__LINE__) + ". Set cuda threads.");
+            LOG(LogLevel::ERROR, "Set cuda threads.");
         }
     }
 }
diff --git a/src/self_attention_cpu.cpp b/src/self_attention_cpu.cpp
index dc74801fd..ed669552d 100644
--- a/src/self_attention_cpu.cpp
+++ b/src/self_attention_cpu.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         self_attention_cpu.cpp
-// Description:  CPU version for self attention
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      March 13, 2023
-// Updated:      June 09, 2023
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
-
 #include "../include/self_attention_cpu.h"
 std::vector<float> transpose(std::vector<float> input,
                              std::vector<int> input_shape,
diff --git a/src/sequential.cpp b/src/sequential.cpp
index 313f77a38..b7d4eaa71 100644
--- a/src/sequential.cpp
+++ b/src/sequential.cpp
@@ -1,19 +1,10 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         sequential.cpp
-// Description:  ...
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      October 09, 2023
-// Updated:      July 19, 2024
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-////////////////////////////////////////////////////////////////////////////////
 
 #include "../include/sequential.h"
 
 #include "../include/config.h"
 #include "../include/conv2d_layer.h"
+#include "../include/custom_logger.h"
 #include "../include/pooling_layer.h"
-// #include "slinear_layer.h"
 #ifdef USE_CUDA
 #include "../include/base_layer_cuda.cuh"
 #endif
@@ -91,9 +82,7 @@ it will be corrected at the first run in the forward pass.
     } else if (this->device.compare("cuda") == 0) {
         this->layers.push_back(layer->to_cuda());
     } else {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Invalid device: [" + this->device + "]");
+        LOG(LogLevel::ERROR, "Invalid device: [" + this->device + "]");
     }
 }
 
@@ -164,17 +153,12 @@ void Sequential::init_output_state_buffer()
             this->temp_states = std::make_shared<TempStateCuda>(
                 this->z_buffer_size, this->z_buffer_block_size);
         } else {
-            throw std::invalid_argument(
-                "Error in file: " + std::string(__FILE__) +
-                " at line: " + std::to_string(__LINE__) +
-                ". Smoothing feature does not support CUDA");
+            LOG(LogLevel::ERROR, "Smoothing feature does not support CUDA");
         }
     }
 #endif
     else {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Invalid device: [" + this->device + "]");
+        LOG(LogLevel::ERROR, "Invalid device: [" + this->device + "]");
     }
 }
 
@@ -197,9 +181,7 @@ void Sequential::init_delta_state_buffer()
     }
 #endif
     else {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Invalid device: [" + this->device + "]");
+        LOG(LogLevel::ERROR, "Invalid device: [" + this->device + "]");
     }
 }
 
@@ -249,8 +231,15 @@ void Sequential::forward(const std::vector<float> &mu_x,
 /*
  */
 {
-    // Batch size
-    int batch_size = mu_x.size() / this->layers.front()->get_input_size();
+    // Batch size: TODO: this is only correct if input size is correctly set
+    int input_size = this->layers.front()->get_input_size();
+    if (mu_x.size() % input_size != 0) {
+        std::string message =
+            "Input size mismatch: " + std::to_string(input_size) + " vs " +
+            std::to_string(mu_x.size());
+        LOG(LogLevel::ERROR, message);
+    }
+    int batch_size = mu_x.size() / input_size;
 
     // Lazy initialization
     if (this->z_buffer_block_size == 0) {
@@ -690,9 +679,7 @@ void Sequential::load_state_dict(
 
 void Sequential::params_from(const Sequential &model_ref) {
     if (this->layers.size() != model_ref.layers.size()) {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". Model architecture is different");
+        LOG(LogLevel::ERROR, "Model architecture is different.");
     }
 
     // TODO: need to add more checks before copying
@@ -780,9 +767,7 @@ std::tuple<pybind11::array_t<float>, pybind11::array_t<float>>
 Sequential::get_input_states() {
     // Check if input_state_update is enabled
     if (!this->input_state_update) {
-        throw std::invalid_argument("Error in file: " + std::string(__FILE__) +
-                                    " at line: " + std::to_string(__LINE__) +
-                                    ". input_state_update is set to False");
+        LOG(LogLevel::ERROR, "input_state_update is set to False.");
     }
 
 #ifdef USE_CUDA
diff --git a/src/slinear_layer.cpp b/src/slinear_layer.cpp
index 9604ff624..d1574fbfb 100644
--- a/src/slinear_layer.cpp
+++ b/src/slinear_layer.cpp
@@ -1,17 +1,7 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         slinear_layer.h
-// Description:  Linear layer which has smoother function. It is used only as
-// the last layer of stacked LSTM networks (SLSTM).
-// Authors:     Van -Dai Vuong, Luong-Ha Nguyen & James-A. Goulet
-// Created:     August 21, 2024
-// Updated:     August 21, 2024
-// Contact:     van-dai.vuong@polymtl.ca, luongha.nguyen@gmail.com &
-// james.goulet@polymtl.ca
-// License:      This code is released under the MIT
-////////////////////////////////////////////////////////////////////////////////
 #include "../include/slinear_layer.h"
 
 #include "../include/common.h"
+#include "../include/custom_logger.h"
 #include "../include/linear_layer.h"
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -106,6 +96,14 @@ void SLinear::forward(BaseHiddenStates &input_states,
 /*
  */
 {
+    // Checkout input size
+    if (this->input_size != input_states.actual_size) {
+        std::string message =
+            "Input size mismatch: " + std::to_string(this->input_size) +
+            " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     // New poitner will point to the same memory location when casting
     SmoothingHiddenStates *smooth_input_states =
         dynamic_cast<SmoothingHiddenStates *>(&input_states);
diff --git a/src/slstm_layer.cpp b/src/slstm_layer.cpp
index e10bc9d74..ecaeea587 100644
--- a/src/slstm_layer.cpp
+++ b/src/slstm_layer.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         slstm_layer.cpp
-// Description:  Long-Short Term Memory (LSTM) which has smoother function in
-// TAGI Authors: Van-Dai Vuong, Luong-Ha Nguyen & James-A. Goulet
-// Created: August 21, 2024
-// Updated : August 21, 2024
-// Contact : van-dai.vuong@polymtl.ca, luongha.nguyen@gmail.com &
-// james.goulet@polymtl.ca
-// License:  This code is released under the MIT
-////////////////////////////////////////////////////////////////////////////////
 #include "slstm_layer.h"
 
 #include <cmath>
@@ -15,6 +5,7 @@
 
 #include "../include/activation.h"
 #include "../include/common.h"
+#include "../include/custom_logger.h"
 #include "../include/lstm_layer.h"
 #include "../include/param_init.h"
 
@@ -225,6 +216,14 @@ void SLSTM::forward(BaseHiddenStates &input_states,
 /*
  */
 {
+    // Checkout input size
+    if (this->input_size * this->seq_len != input_states.actual_size) {
+        std::string message = "Input size mismatch: " +
+                              std::to_string(this->input_size * this->seq_len) +
+                              " vs " + std::to_string(input_states.actual_size);
+        LOG(LogLevel::ERROR, message);
+    }
+
     // New poitner will point to the same memory location when casting
     SmoothingHiddenStates *smooth_input_states =
         dynamic_cast<SmoothingHiddenStates *>(&input_states);
diff --git a/src/utils.cpp b/src/utils.cpp
index cc368403c..079f947de 100644
--- a/src/utils.cpp
+++ b/src/utils.cpp
@@ -1,13 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         utils.cpp
-// Description:  utils for TAGI package such saving and loading parameters
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      April 10, 2022
-// Updated:      July 20, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
-
 #include "../include/utils.h"
 
 void save_error_rate(std::string &res_path, std::vector<float> &error_rate,


Test patch:
diff --git a/include/test_unit.h b/include/test_unit.h
index 62d899f59..f9f670b7f 100644
--- a/include/test_unit.h
+++ b/include/test_unit.h
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         test_units.h
-// Description:  Header file for test units
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      August 21, 2022
-// Updated:      August 21, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
 #pragma once
 #include <iomanip>
 #include <iostream>
diff --git a/src/test_dataloader.cpp b/src/test_dataloader.cpp
index b4afc5c54..84584476b 100644
--- a/src/test_dataloader.cpp
+++ b/src/test_dataloader.cpp
@@ -1,12 +1,3 @@
-///////////////////////////////////////////////////////////////////////////////
-// File:         test_dataloader.cpp
-// Description:  Testing unit for dataloader
-// Authors:      Luong-Ha Nguyen & James-A. Goulet
-// Created:      August 21, 2022
-// Updated:      August 21, 2022
-// Contact:      luongha.nguyen@gmail.com & james.goulet@polymtl.ca
-// License:      This code is released under the MIT License.
-///////////////////////////////////////////////////////////////////////////////
 #include "../include/test_unit.h"
 int main() {
     // Data
diff --git a/test/cpp_unit/test_mnist.cpp b/test/cpp_unit/test_mnist.cpp
index be3909143..ad2860995 100644
--- a/test/cpp_unit/test_mnist.cpp
+++ b/test/cpp_unit/test_mnist.cpp
@@ -359,6 +359,14 @@ TEST_F(MnistTest, LayerNormCNNTest_CPU) {
     EXPECT_LT(avg_error, threshold) << "Error rate is higher than threshold";
 }
 
+TEST_F(MnistTest, MismatchSizeDetection) {
+    Sequential model(Linear(784, 32), ReLU(), Linear(300, 32), ReLU(),
+                     Linear(32, 11));
+    float avg_error;
+    float threshold = 0.5;
+    EXPECT_THROW({ mnist_test_runner(model, avg_error); }, std::runtime_error);
+}
+
 #ifdef USE_CUDA
 TEST_F(MnistTest, FNNModelTest_CUDA) {
     if (!g_gpu_enabled) GTEST_SKIP() << "GPU tests are disabled.";
diff --git a/test/fnn/test_fnn_mnist_cpu.cpp b/test/fnn/test_fnn_mnist_cpu.cpp
index c409cb7a6..5f62481ca 100644
--- a/test/fnn/test_fnn_mnist_cpu.cpp
+++ b/test/fnn/test_fnn_mnist_cpu.cpp
@@ -66,17 +66,17 @@ void fnn_mnist() {
     //////////////////////////////////////////////////////////////////////
     // TAGI network
     //////////////////////////////////////////////////////////////////////
-    Sequential model(Linear(784, 16), ReLU(), Linear(16, 16), ReLU(),
-                     Linear(16, 11));
+    // Sequential model(Linear(784, 16), ReLU(), Linear(32, 16), ReLU(),
+    //                  Linear(16, 11));
 
     // Sequential model(Linear(784, 1024), BatchNorm2d(1024), ReLU(),
     //                  Linear(1024, 1024), BatchNorm2d(1024), ReLU(),
     //                  Linear(1024, 11));
 
-    // Sequential model(Linear(784, 100), LayerNorm(std::vector<int>({100})),
-    //                  ReLU(), Linear(100, 100),
-    //                  LayerNorm(std::vector<int>({100})), ReLU(),
-    //                  Linear(100, 11));
+    Sequential model(Linear(784, 100), LayerNorm(std::vector<int>({100})),
+                     ReLU(), Linear(100, 100),
+                     LayerNorm(std::vector<int>({100})), ReLU(),
+                     Linear(100, 11));
 
     // Sequential model(Conv2d(1, 16, 4, true, 1, 1, 1, 28, 28), ReLU(),
     //                  AvgPool2d(3, 2), Conv2d(16, 32, 5), ReLU(),