Repository: kovanostra/message-passing-neural-network (language: python)

Issue:
Create LICENSE

Refactor training
Introduces ModelTrainer
Combines graphencoder and fully connected layer into a single model
Configure pytorch
Changes algorithm from numpy implementation to pytorch


Discussion / hints:

Fix patch:
diff --git a/Dockerfile b/Dockerfile
deleted file mode 100644
index 9294111..0000000
--- a/Dockerfile
+++ /dev/null
@@ -1,15 +0,0 @@
-FROM python:3.7.6
-
-RUN mkdir /app/
-
-COPY /dist/message-passing-nn-*tar.gz /app/app.tar.gz
-COPY /data /data
-COPY ./docker-entrypoint.sh /app/
-COPY parameters/grid-search-parameters.sh /app/
-RUN chmod 777 /app/docker-entrypoint.sh && ln -s /app/docker-entrypoint.sh /
-
-RUN pip install --upgrade pip && pip install /app/app.tar.gz
-
-ENTRYPOINT ["sh", "/app/docker-entrypoint.sh"]
-RUN cd /app/
-CMD ["message-passing-nn", "grid-search"]
\ No newline at end of file
diff --git a/ENVIRONMENT_VARIABLES.md b/ENVIRONMENT_VARIABLES.md
index ed40535..1a39fa8 100644
--- a/ENVIRONMENT_VARIABLES.md
+++ b/ENVIRONMENT_VARIABLES.md
@@ -16,21 +16,17 @@ MODEL_DIRECTORY='model_checkpoints'
 
 RESULTS_DIRECTORY='grid_search_results'
 
-- The option to run the model on 'cpu' or 'cuda' can be controlled by (*cuda recommended only for the 'RNN' model*): 
+- The option to run the model on 'cpu' or 'cuda' can be controlled by: 
 
 DEVICE='cpu'
 
-- If the model runs on 'cpu', then the option to use multiprocessing (will use ALL available cores) or not comes from:
-
-CPU_MULTIPROCESSING='True' or 'False'
-
 **USED FOR GRID SEARCH**
 
 To define a range for the grid search please pass the values in the following format:
 1. For numeric ranges: ENVVAR='min_value&max_value&number_of_values' (e.g. '10&15&2')
 2. For string ranges: ENVVAR='selection_1&selection_2' (e.g. 'SGD&Adam')
 
-- The model to use ('RNN' or 'GRU') is defined by :
+- The model to use (only 'RNN' available at this version') is defined by :
 
 MODEL='RNN'
 
@@ -50,14 +46,6 @@ OPTIMIZER='SGD'
 
 BATCH_SIZE='1'
 
-- If the number of features in your dataset is too large please change the following value (-1 will use all of them)
-
-MAXIMUM_NUMBER_OF_FEATURES='-1'
-
-- If the number of nodes in your dataset is too large please change the following value (-1 will use all of them)
-
-MAXIMUM_NUMBER_OF_NODES='-1'
-
 - The validation split can be controlled by:
 
 VALIDATION_SPLIT='0.2'
@@ -76,6 +64,6 @@ VALIDATION_PERIOD='5'
 
 **USED FOR INFERENCE**
 
-- The model to load ('RNN' or 'GRU') is defined by :
+- The model to load (only 'RNN' available at this version') is defined by :
 
 MODEL='RNN'
\ No newline at end of file
diff --git a/README.md b/README.md
index 0e027e2..f26da94 100644
--- a/README.md
+++ b/README.md
@@ -1,54 +1,67 @@
-[![Build Status](https://dev.azure.com/kovamos/message-passing-nn/_apis/build/status/kovanostra.message-passing-nn?branchName=master)](https://dev.azure.com/kovamos/message-passing-nn/_build/latest?definitionId=2&branchName=master)
 ### Table of contents
 - [1. Description](#1-description)
-- [2. Build and use](#2-build-and-use)
-- [3. Import and use](#3-import-as-package)
-- [4. Examples](#4-examples)
-- [5. Requirements](#5-requirements)
-- [6. Environment](#6-environment)
-- [7. Dataset](#7-dataset)
-- [8. Environment variables](#8-environment-variables)
-- [9. Execute a grid search](#9-execute-a-grid-search)
-- [10. Tox build](#10-tox-build)
-- [11. Run the code using docker](#11-run-the-code-using-docker)
-- [12. Azure pipelines](#12-azure-pipelines)
-- [13. Upcoming](#13-upcoming)
+- [2. Model architecture](#2-model-architecture)
+- [3. Build and use](#3-build-and-use)
+- [4. Import and use](#4-import-as-package)
+- [5. Examples](#5-examples)
+- [6. Requirements](#6-requirements)
+- [7. Environment](#7-environment)
+- [8. Dataset](#8-dataset)
+- [9. Environment variables](#9-environment-variables)
+- [10. Execute a grid search](#10-execute-a-grid-search)
+- [11. Execute an inference](#11-execute-an-inference)
 
 
 ### 1. Description
 
 This repository contains:
-1. A pytorch implementation of a message passing neural network with either RNN or GRU units (inspired from https://arxiv.org/abs/1812.01070). 
-2. A wrapper around the model to perform a grid search, and save model checkpoints when the validation error is best for each configuration.
+1. A pytorch C++ implementation of a message passing neural network with RNN units (inspired from https://arxiv.org/abs/1812.01070). 
+2. A python wrapper around the model to perform a grid search, and save model checkpoints for each validation step.
+3. A script to perform an inference on a dataset based on a specific model checkpoint.
+4. A custom CUDA kernel.
 
-### 2. Build and use
+### 2. Model Architecture
 
-To use the current version (1.5.0) you need to first build the project. Please clone the repository and then run the build scripts depending on your OS.
+![Model Architecture](./forward_pass_equations.png)
 
-Linux
+### 3. Build and use
+
+To use the current version (master or tag >= 1.5.0) you need to first build the project. Please clone the repository and then run the build scripts depending on your OS and whether you have a CUDA enabled GPU available. To build for GPU you first need to set the CUDA_HOME variable in the respective .sh file (here it defaults to usr/local/cuda).
+
+Linux (CPU & GPU)
 ```
 . linux_build.sh
 ```
 
-macOS
+Linux (CPU only)
+```
+. linux_build_cpu.sh
+```
+
+macOS (CPU & GPU)
 ```
 . macos_build.sh
 ```
 
-Then you can use the code as in the [examples](#4-examples) or perform a [grid search](#9-execute-a-grid-search). The GRU model can be used without building as it is only python code.
+macOS (CPU only)
+```
+. macos_build_cpu.sh
+```
+
+Then you can use the code as in the [examples](#4-examples) or perform a [grid search](#9-execute-a-grid-search). 
 
-### 3. Import as package
-If you can't build the project you can install the pure python version of the project (version 1.4.2) using pip please run:
+### 4. Import as package
+If you can't build the project you can install the pure python version of the project (version 1.4.2) using pip:
 
 ```
 pip install message-passing-nn
 ```
 
-### 4. Examples
+### 5. Examples
 
-The code can be used to either train a single configuration of the message passing neural network or to perform a grid search. For usage examples please look in the example_notebooks/ directory or on the [colab notebook](https://colab.research.google.com/drive/1jFJ7l7jIv22BhvvzlmXOWFtgBE15ea2X).
+The code can be used to either train a single configuration of the message passing neural network or to perform a grid search. For usage examples (v1.4.2) please look in the example_notebooks/ directory or on the [colab notebook](https://colab.research.google.com/drive/1jFJ7l7jIv22BhvvzlmXOWFtgBE15ea2X).
 
-### 5. Requirements
+### 6. Requirements
 
 Python 3.7.6
 
@@ -58,6 +71,7 @@ click
 torch=1.5.0
 numpy==1.17.4
 pandas=1.0.3
+tqdm
 ```
 
 Tests
@@ -67,14 +81,14 @@ torch=1.5.0
 pandas=1.0.3
 ```
 
-### 6. Environment
-To create the "message-passing-nn" conda environment please run:
+### 7. Environment
+To create the "message-passing-neural-network" conda environment please run:
 
 ```
 conda env create -f environment.yml
 ```
 
-### 7. Dataset
+### 8. Dataset
       
 The repository expects the data to be in the following format:
 
@@ -96,11 +110,11 @@ This repository contains two dataset folders with examples of data to run the co
   - sample-dataset: Contains just one pair of features/labels with some default values. This data lets you run the code in demo mode.
   - protein-folding: Contains pairs of features/labels for various proteins (prepared using https://github.com/simonholmes001/structure_prediction). The features represent protein characteristics, and the labels the distance between all aminoacids.
 
-### 8. Environment variables
+### 9. Environment variables
       
 The model and grid search can be set up using a set of environment variables contained in the grid-search-parameters.sh. Please refer to the ENVIRONMENT_VARIABLES.md for the full list of available environment variables and how to use them.
 
-### 9. Execute a grid search
+### 10. Execute a grid search
        
 Before executing a grid-search please go to the grid-search.sh to add your PYTHONPATH=path/to/message-passing-nn/.
 
@@ -113,57 +127,26 @@ This script will:
 
 1. Create the conda environment from the environment.yml (if not created already)
 2. Activate it
-3. It exports the PYTHONPATH=path/to/message-passing-nn/ (line needs to be uncommented first)
+3. Export the PYTHONPATH=path/to/message-passing-nn/ (line needs to be uncommented first)
 4. Export the environment variables to be used for the Grid Search
 5. Run the grid search
+6. Save model checkpoints for each validation and a csv file containing all calculated losses
 
-### 10. Tox build
+### 11. Execute an inference
 
-Tox is a tool which downloads the code dependencies, runs all the tests and, if the tests pass, it builds an artifact in the .tox/dist/ directory. The artifact is name tagged by the version of your code which can be specified in the setup.py.
+Before executing an inference please go to the inference.sh to add your PYTHONPATH=path/to/message-passing-nn/. Please also make sure that the dataset used for inference is of same dimensions (M, N, L) as the one used to train the model.
 
-Requirements
-
-```
-click
-tox==3.14.3
-pytorch=1.4.0
-numpy==1.17.4
-pandas=1.03
-```
-
-From the parent directory run (with sudo if necessary):
-```
-tox
-```
-
-### 11. Run the code using docker
-The model can be run from inside a docker container (currently cpu only). To do so please execute the following shell script:
+The grid search can be executed by executing a shell script:
 ```
-. docker-grid-search.sh
+. grid-search.sh
 ```
 
-The grid-search-docker.sh will:
-
-    1. Remove any previous message-passing-nn containers and images
-    2. Build the project
-    3. Create a docker image
-    4. Create a docker container
-    5. Start the container
-    6. Print the containner's logs with the --follow option activated
-
-By default the dockerfile uses the sample-dataset. To change that please modify the grid-search-parameters.sh.
-
-You can clear the docker container and images created by running:
-```
-. remove-containers-and-images.sh
-```
-This, by default will remove only tagged images created by the grid-search-docker.sh. However, you can uncomment the following lines if you want to remove all stopped containers and untagged images too:
-```
-docker container rm $(docker container ls -aq)
-docker rmi $(docker images | grep "^<none>" | awk "{print $3}")
-```
-Please note that this will delete also untagged images created by other projects, so use with caution.
+This script will:
 
-### 12. Azure pipelines
+1. Create the conda environment from the environment.yml (if not created already)
+2. Activate it
+3. Export the PYTHONPATH=path/to/message-passing-nn/ (line needs to be uncommented first)
+4. Export the environment variables to be used for the Inference
+5. Run the inference
+6. Save results as a list of (output, label, tag) for each input
 
-https://dev.azure.com/kovamos/message-passing-nn
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
deleted file mode 100644
index 49e77ad..0000000
--- a/azure-pipelines.yml
+++ /dev/null
@@ -1,26 +0,0 @@
-# Python package
-# Create and test a Python package on multiple Python versions.
-# Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more:
-# https://docs.microsoft.com/azure/devops/pipelines/languages/python
-
-trigger:
-  branches:
-    include:
-      - '*'
-
-pool:
-  vmImage: 'ubuntu-latest'
-strategy:
-  matrix:
-    Python377:
-      python.version: '3.7.7'
-
-steps:
-- task: UsePythonVersion@0
-  inputs:
-    versionSpec: '$(python.version)'
-  displayName: 'Use Python $(python.version)'
-
-- script: |
-    python -m unittest discover
-  displayName: 'Run the tests'
diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh
deleted file mode 100644
index 9ead2f7..0000000
--- a/docker-entrypoint.sh
+++ /dev/null
@@ -1,4 +0,0 @@
-#!/bin/bash
-
- . "/app/grid-search-parameters.sh"
-exec "$@"
\ No newline at end of file
diff --git a/docker-grid-search.sh b/docker-grid-search.sh
deleted file mode 100644
index 2d9a0d7..0000000
--- a/docker-grid-search.sh
+++ /dev/null
@@ -1,16 +0,0 @@
-#!/bin/bash
-
-python setup.py sdist
-
-echo "Removing previous containers and images for message-passing-nn"
-docker rm -f message-passing-nn-container
-docker rmi message-passing-nn
-
-echo "Building a new docker image and container"
-docker build . -t message-passing-nn
-docker container create --name message-passing-nn-container message-passing-nn:latest
-
-echo "Starting message-passing-nn-container"
-docker start message-passing-nn-container
-
-docker logs -f message-passing-nn-container
diff --git a/environment.yml b/environment.yml
index 98f6dad..120efd9 100644
--- a/environment.yml
+++ b/environment.yml
@@ -10,4 +10,5 @@ dependencies:
   - pandas=1.0.3
   - pytorch=1.5.0
   - setuptools=46.0.0
-  - click
\ No newline at end of file
+  - click
+  - tqdm
\ No newline at end of file
diff --git a/example_notebooks/Train a single configuration.ipynb b/example_notebooks/Train a single configuration.ipynb
index 639d268..dda84fb 100644
--- a/example_notebooks/Train a single configuration.ipynb
+++ b/example_notebooks/Train a single configuration.ipynb
@@ -21,7 +21,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Clone the repository to get the data folders"
+    "## Clone the infrastructure to get the data folders"
    ]
   },
   {
@@ -52,7 +52,7 @@
     "from message_passing_nn.graph.graph_rnn_encoder import GraphRNNEncoder\n",
     "from message_passing_nn.graph.graph_gru_encoder import GraphGRUEncoder\n",
     "from message_passing_nn.data.data_preprocessor import DataPreprocessor\n",
-    "from message_passing_nn.repository.file_system_repository import FileSystemRepository"
+    "from message_passing_nn.infrastructure.file_system_repository import FileSystemRepository"
    ]
   },
   {
diff --git a/forward_pass_equations.png b/forward_pass_equations.png
new file mode 100644
index 0000000..483dafa
Binary files /dev/null and b/forward_pass_equations.png differ
diff --git a/linux_build.sh b/linux_build.sh
index c4c42a9..2fb9058 100644
--- a/linux_build.sh
+++ b/linux_build.sh
@@ -1,5 +1,6 @@
 #!/bin/bash
 
+export CUDA_HOME=/usr/local/cuda
 export CC=gcc
 export CXX=g++
 python setup.py install
\ No newline at end of file
diff --git a/linux_build_cpu.sh b/linux_build_cpu.sh
new file mode 100644
index 0000000..d8af700
--- /dev/null
+++ b/linux_build_cpu.sh
@@ -0,0 +1,5 @@
+#!/bin/bash
+
+export CC=gcc
+export CXX=g++
+python setup_cpu.py install
\ No newline at end of file
diff --git a/macos_build.sh b/macos_build.sh
index d149a12..14ef91b 100644
--- a/macos_build.sh
+++ b/macos_build.sh
@@ -1,5 +1,6 @@
 #!/bin/bash
 
+export CUDA_HOME=/usr/local/cuda
 export MACOSX_DEPLOYMENT_TARGET=10.11
 export CC=clang
 export CXX=clang++
diff --git a/macos_build_cpu.sh b/macos_build_cpu.sh
new file mode 100644
index 0000000..092b54e
--- /dev/null
+++ b/macos_build_cpu.sh
@@ -0,0 +1,6 @@
+#!/bin/bash
+
+export MACOSX_DEPLOYMENT_TARGET=10.11
+export CC=clang
+export CXX=clang++
+python setup_cpu.py install
diff --git a/message_passing_nn/cli.py b/message_passing_nn/cli.py
index 409963e..a7d31db 100644
--- a/message_passing_nn/cli.py
+++ b/message_passing_nn/cli.py
@@ -20,13 +20,10 @@ def main(debug):
 @click.argument('results_directory', envvar='RESULTS_DIRECTORY', type=str)
 @click.argument('model', envvar='MODEL', type=str)
 @click.argument('device', envvar='DEVICE', type=str)
-@click.argument('cpu_multiprocessing', envvar='CPU_MULTIPROCESSING', type=str)
 @click.argument('epochs', envvar='EPOCHS', type=str)
 @click.argument('loss_function', envvar='LOSS_FUNCTION', type=str)
 @click.argument('optimizer', envvar='OPTIMIZER', type=str)
 @click.argument('batch_size', envvar='BATCH_SIZE', type=str)
-@click.argument('maximum_number_of_features', envvar='MAXIMUM_NUMBER_OF_FEATURES', type=str)
-@click.argument('maximum_number_of_nodes', envvar='MAXIMUM_NUMBER_OF_NODES', type=str)
 @click.argument('validation_split', envvar='VALIDATION_SPLIT', type=str)
 @click.argument('test_split', envvar='TEST_SPLIT', type=str)
 @click.argument('time_steps', envvar='TIME_STEPS', type=str)
@@ -36,32 +33,25 @@ def start_training(dataset_name: str,
                    model_directory: str,
                    results_directory: str,
                    model: str,
-                   cpu_multiprocessing: str,
                    device: str,
                    epochs: str,
                    loss_function: str,
                    optimizer: str,
                    batch_size: str,
-                   maximum_number_of_features: str,
-                   maximum_number_of_nodes: str,
                    validation_split: str,
                    test_split: str,
                    time_steps: str,
                    validation_period: str) -> None:
-    get_logger().info("Starting training")
     message_passing_nn = create_grid_search(dataset_name,
                                             data_directory,
                                             model_directory,
                                             results_directory,
                                             model,
                                             device,
-                                            cpu_multiprocessing,
                                             epochs,
                                             loss_function,
                                             optimizer,
                                             batch_size,
-                                            maximum_number_of_features,
-                                            maximum_number_of_nodes,
                                             validation_split,
                                             test_split,
                                             time_steps,
diff --git a/message_passing_nn/create_message_passing_nn.py b/message_passing_nn/create_message_passing_nn.py
index eaa9ab2..5c79f7d 100644
--- a/message_passing_nn/create_message_passing_nn.py
+++ b/message_passing_nn/create_message_passing_nn.py
@@ -4,7 +4,6 @@
 from message_passing_nn.model.inferencer import Inferencer
 from message_passing_nn.model.loader import Loader
 from message_passing_nn.model.trainer import Trainer
-from message_passing_nn.repository.file_system_repository import FileSystemRepository
 from message_passing_nn.usecase import Usecase
 from message_passing_nn.usecase.grid_search import GridSearch
 from message_passing_nn.usecase.inference import Inference
@@ -29,13 +28,10 @@ def create_grid_search(dataset_name: str,
                        results_directory: str,
                        model: str,
                        device: str,
-                       cpu_multiprocessing: str,
                        epochs: str,
                        loss_function_selection: str,
                        optimizer_selection: str,
                        batch_size: str,
-                       maximum_number_of_features: str,
-                       maximum_number_of_nodes: str,
                        validation_split: str,
                        test_split: str,
                        time_steps: str,
@@ -45,17 +41,19 @@ def create_grid_search(dataset_name: str,
                                                                                      loss_function_selection,
                                                                                      optimizer_selection,
                                                                                      batch_size,
-                                                                                     maximum_number_of_features,
-                                                                                     maximum_number_of_nodes,
                                                                                      validation_split,
                                                                                      test_split,
                                                                                      time_steps,
                                                                                      validation_period)
-    file_system_repository = FileSystemRepository(data_directory, dataset_name)
+    data_path = _get_data_path(data_directory, dataset_name)
     data_preprocessor = DataPreprocessor()
-    trainer = Trainer(data_preprocessor, device, eval(cpu_multiprocessing.capitalize()))
+    trainer = Trainer(data_preprocessor, device)
     saver = Saver(model_directory, results_directory)
-    grid_search = GridSearch(file_system_repository, data_preprocessor, trainer, grid_search_dictionary, saver)
+    grid_search = GridSearch(data_path,
+                             data_preprocessor,
+                             trainer,
+                             grid_search_dictionary,
+                             saver)
     return MessagePassingNN(grid_search)
 
 
@@ -65,14 +63,18 @@ def create_inference(dataset_name: str,
                      results_directory: str,
                      model: str,
                      device: str) -> MessagePassingNN:
-    file_system_repository = FileSystemRepository(data_directory, dataset_name)
+    data_path = data_directory + dataset_name + "/"
     data_preprocessor = DataPreprocessor()
     model_loader = Loader(model)
     model_inferencer = Inferencer(data_preprocessor, device)
     saver = Saver(model_directory, results_directory)
-    inference = Inference(file_system_repository, data_preprocessor, model_loader, model_inferencer, saver)
+    inference = Inference(data_path, data_preprocessor, model_loader, model_inferencer, saver)
     return MessagePassingNN(inference)
 
 
+def _get_data_path(data_directory: str, dataset_name: str) -> str:
+    return data_directory + dataset_name + "/"
+
+
 def get_logger() -> logging.Logger:
     return logging.getLogger('message_passing_nn')
diff --git a/message_passing_nn/data/__init__.py b/message_passing_nn/data/__init__.py
index e89a17f..bf95ea7 100644
--- a/message_passing_nn/data/__init__.py
+++ b/message_passing_nn/data/__init__.py
@@ -1,3 +1,2 @@
 from message_passing_nn.data.data_preprocessor import DataPreprocessor
-from message_passing_nn.data.graph_dataset import GraphDataset
 from message_passing_nn.data.preprocessor import Preprocessor
diff --git a/message_passing_nn/data/data_preprocessor.py b/message_passing_nn/data/data_preprocessor.py
index fef5101..0ab9fbe 100644
--- a/message_passing_nn/data/data_preprocessor.py
+++ b/message_passing_nn/data/data_preprocessor.py
@@ -4,30 +4,37 @@
 import torch as to
 from torch import nn
 from torch.utils.data import DataLoader
+from torch.utils.data.sampler import SubsetRandomSampler
+from tqdm import tqdm
 
-from message_passing_nn.data.graph_dataset import GraphDataset
+from message_passing_nn.infrastructure.graph_dataset import GraphDataset
 from message_passing_nn.data.preprocessor import Preprocessor
 
 
 class DataPreprocessor(Preprocessor):
     def __init__(self):
         super().__init__()
+        self.test_mode = False
 
     def train_validation_test_split(self,
-                                    dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]],
+                                    dataset: GraphDataset,
                                     batch_size: int,
                                     validation_split: float = 0.2,
                                     test_split: float = 0.1) -> Tuple[DataLoader, DataLoader, DataLoader]:
         test_index, validation_index = self._get_validation_and_test_indexes(dataset,
                                                                              validation_split,
                                                                              test_split)
-        training_data = DataLoader(GraphDataset(dataset[:validation_index]), batch_size)
+        train_sampler = SubsetRandomSampler(list(range(validation_index)))
+        validation_sampler = SubsetRandomSampler(list(range(validation_index, test_index)))
+        test_sampler = SubsetRandomSampler(list(range(test_index, len(dataset.dataset))))
+
+        training_data = DataLoader(dataset, batch_size=batch_size, sampler=train_sampler)
         if validation_split:
-            validation_data = DataLoader(GraphDataset(dataset[validation_index:test_index]), batch_size)
+            validation_data = DataLoader(dataset, batch_size=batch_size, sampler=validation_sampler)
         else:
             validation_data = DataLoader(GraphDataset([]))
         if test_split:
-            test_data = DataLoader(GraphDataset(dataset[test_index:]), batch_size)
+            test_data = DataLoader(dataset, batch_size=batch_size, sampler=test_sampler)
         else:
             test_data = DataLoader(GraphDataset([]))
         self.get_logger().info("Train/validation/test split: " + "/".join([str(len(training_data)),
@@ -37,32 +44,29 @@ def train_validation_test_split(self,
         return training_data, validation_data, test_data
 
     @staticmethod
-    def get_dataloader(dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]], batch_size: int = 1) -> DataLoader:
-        return DataLoader(GraphDataset(dataset), batch_size)
-
-    def equalize_dataset_dimensions(self,
-                                    dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]],
-                                    maximum_number_of_nodes: int = -1,
-                                    maximum_number_of_features: int = -1) \
-            -> List[Tuple[to.Tensor, to.Tensor, to.Tensor]]:
-
-        adjacency_matrix_max_size, features_max_size, labels_max_size = self._get_maximum_data_size(dataset,
-                                                                                                    maximum_number_of_nodes)
-
-        preprocessed_dataset = self._equalize_sizes(dataset,
-                                                    adjacency_matrix_max_size,
-                                                    features_max_size,
-                                                    labels_max_size,
-                                                    maximum_number_of_nodes,
-                                                    maximum_number_of_features)
-        return preprocessed_dataset
+    def get_dataloader(dataset: GraphDataset, batch_size: int = 1) -> DataLoader:
+        return DataLoader(dataset, batch_size)
+
+    def find_all_node_neighbors(self, dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]]) -> List[
+        Tuple[to.Tensor, to.Tensor, to.Tensor]]:
+        dataset_with_neighbors = []
+        disable_progress_bar = self.test_mode
+        for index in tqdm(range(len(dataset)), disable=disable_progress_bar):
+            features, adjacency_matrix, labels = dataset[index]
+            number_of_nodes = features.shape[0]
+            all_neighbors = to.zeros(number_of_nodes, number_of_nodes) - to.ones(number_of_nodes, number_of_nodes)
+            all_neighbors_list = [to.nonzero(adjacency_matrix[node_id], as_tuple=True)[0].tolist() for node_id in
+                                  range(adjacency_matrix.shape[0])]
+            for node_id in range(number_of_nodes):
+                all_neighbors[node_id, :len(all_neighbors_list[node_id])] = to.tensor(all_neighbors_list[node_id])
+            dataset_with_neighbors.append((features, all_neighbors, labels))
+        return dataset_with_neighbors
 
     @staticmethod
-    def extract_data_dimensions(dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]]) -> Tuple:
+    def extract_data_dimensions(dataset: GraphDataset) -> Tuple:
         node_features_size = dataset[0][0].size()
-        adjacency_matrix_size = dataset[0][1].size()
         labels_size = dataset[0][2].size()
-        return node_features_size, adjacency_matrix_size, labels_size
+        return node_features_size, labels_size
 
     @staticmethod
     def flatten(tensors: to.Tensor, desired_size: int = 0) -> to.Tensor:
@@ -86,60 +90,15 @@ def _pad_zeros(flattened_tensor: to.Tensor, desired_size: int) -> to.Tensor:
         return flattened_tensor
 
     @staticmethod
-    def _get_validation_and_test_indexes(dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]],
+    def _get_validation_and_test_indexes(dataset: GraphDataset,
                                          validation_split: float,
                                          test_split: float) -> Tuple[int, int]:
         validation_index = int((1 - validation_split - test_split) * len(dataset))
         test_index = int((1 - test_split) * len(dataset))
         return test_index, validation_index
 
-    def _equalize_sizes(self,
-                        dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]],
-                        adjacency_matrix_max_size: List[int],
-                        features_max_size: List[int],
-                        labels_max_size: List[int],
-                        maximum_number_of_nodes: int,
-                        maximum_number_of_features: int) -> List[Tuple[to.Tensor, to.Tensor, to.Tensor]]:
-        preprocessed_dataset = []
-        for features, adjacency_matrix, labels in dataset:
-            if self._is_graph_size_less_than_maximum_allowed(maximum_number_of_nodes, adjacency_matrix):
-                continue
-            if maximum_number_of_features < 0:
-                maximum_number_of_features = features.size()[1]
-            features_preprocessed = to.zeros((features_max_size[0], maximum_number_of_features))
-            adjacency_matrix_preprocessed = to.zeros((adjacency_matrix_max_size[0], adjacency_matrix_max_size[1]))
-            labels_preprocessed = to.zeros((labels_max_size[0]))
-
-            features_preprocessed[:features.size()[0], :maximum_number_of_features] = features[:,
-                                                                                      :maximum_number_of_features]
-            adjacency_matrix_preprocessed[:adjacency_matrix.size()[0], :adjacency_matrix.size()[1]] = adjacency_matrix
-            labels_preprocessed[:labels.size()[0]] = labels
-            preprocessed_dataset.append((features_preprocessed, adjacency_matrix_preprocessed, labels_preprocessed))
-        return preprocessed_dataset
-
-    def _get_maximum_data_size(self,
-                               dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]],
-                               maximum_number_of_nodes: int) \
-            -> Tuple[List, List, List]:
-        features_max_size = [0, 0]
-        adjacency_matrix_max_size = [0, 0]
-        labels_max_size = [0]
-        for features, adjacency_matrix, labels in dataset:
-            if self._is_graph_size_less_than_maximum_allowed(maximum_number_of_nodes, adjacency_matrix):
-                continue
-            if features.size()[0] > features_max_size[0]:
-                features_max_size[0] = features.size()[0]
-            if features.size()[1] > features_max_size[1]:
-                features_max_size[1] = features.size()[1]
-            if adjacency_matrix.size()[0] > adjacency_matrix_max_size[0]:
-                adjacency_matrix_max_size = adjacency_matrix.size()
-            if labels.size()[0] > labels_max_size[0]:
-                labels_max_size = labels.size()
-        return adjacency_matrix_max_size, features_max_size, labels_max_size
-
-    @staticmethod
-    def _is_graph_size_less_than_maximum_allowed(maximum_number_of_nodes: int, adjacency_matrix: to.Tensor) -> bool:
-        return 0 < maximum_number_of_nodes < adjacency_matrix.size()[0]
+    def enable_test_mode(self) -> None:
+        self.test_mode = True
 
     @staticmethod
     def get_logger() -> logging.Logger:
diff --git a/message_passing_nn/data/graph_dataset.py b/message_passing_nn/data/graph_dataset.py
deleted file mode 100644
index 2ac5ba7..0000000
--- a/message_passing_nn/data/graph_dataset.py
+++ /dev/null
@@ -1,30 +0,0 @@
-from typing import List, Tuple
-
-import torch as to
-from torch.utils.data import Dataset
-
-
-class GraphDataset(Dataset):
-    def __init__(self,
-                 dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]],
-                 normalize_features: bool = True,
-                 normalize_labels: bool = True) -> None:
-        self.features = self._to_list(dataset)
-        self.labels = self._extract_labels(dataset)
-        self.normalize_features = normalize_features
-        self.normalize_labels = normalize_labels
-
-    def __len__(self) -> int:
-        return len(self.features)
-
-    def __getitem__(self, index: int) -> Tuple:
-        return self.features[index], self.labels[index]
-
-    @staticmethod
-    def _to_list(dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]]) -> List[Tuple[to.Tensor, to.Tensor]]:
-        return [(dataset[index][0], dataset[index][1]) for index in range(len(dataset))]
-
-    @staticmethod
-    def _extract_labels(dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]]) -> List[to.Tensor]:
-        return [dataset[index][2] for index in range(len(dataset))]
-
diff --git a/message_passing_nn/graph/__init__.py b/message_passing_nn/graph/__init__.py
index 1a63532..a90b5c0 100644
--- a/message_passing_nn/graph/__init__.py
+++ b/message_passing_nn/graph/__init__.py
@@ -1,3 +1,1 @@
-from message_passing_nn.graph.node import Node
 from message_passing_nn.graph.rnn_encoder import RNNEncoder
-# from message_passing_nn.graph.gru_encoder import GRUEncoder
diff --git a/message_passing_nn/graph/gru_encoder.py b/message_passing_nn/graph/gru_encoder.py
deleted file mode 100644
index f6891bf..0000000
--- a/message_passing_nn/graph/gru_encoder.py
+++ /dev/null
@@ -1,197 +0,0 @@
-from typing import List
-
-import torch as to
-import torch.nn as nn
-
-from message_passing_nn.data.data_preprocessor import DataPreprocessor
-from message_passing_nn.graph.node import Node
-
-
-class GRUEncoder(nn.Module):
-    def __init__(self,
-                 time_steps: int,
-                 number_of_nodes: int,
-                 number_of_node_features: int,
-                 fully_connected_layer_input_size: int,
-                 fully_connected_layer_output_size: int) -> None:
-        super(GRUEncoder, self).__init__()
-        base_tensor_shape = [number_of_node_features, number_of_node_features]
-        base_2d_tensor_shape = [number_of_node_features, 1]
-
-        self.time_steps = time_steps
-        self.number_of_nodes = number_of_nodes
-        self.number_of_node_features = number_of_node_features
-        self.fully_connected_layer_input_size = fully_connected_layer_input_size
-        self.fully_connected_layer_output_size = fully_connected_layer_output_size
-
-        self.w_gru_update_gate_features = self._get_parameter(base_tensor_shape)
-        self.w_gru_forget_gate_features = self._get_parameter(base_tensor_shape)
-        self.w_gru_current_memory_message_features = self._get_parameter(base_tensor_shape)
-        self.u_gru_update_gate = self._get_parameter(base_tensor_shape)
-        self.u_gru_forget_gate = self._get_parameter(base_tensor_shape)
-        self.u_gru_current_memory_message = self._get_parameter(base_tensor_shape)
-        self.b_gru_update_gate = self._get_parameter(base_2d_tensor_shape).view(base_2d_tensor_shape[0])
-        self.b_gru_forget_gate = self._get_parameter(base_2d_tensor_shape).view(base_2d_tensor_shape[0])
-        self.b_gru_current_memory_message = self._get_parameter(base_2d_tensor_shape).view(base_2d_tensor_shape[0])
-        self.u_graph_node_features = self._get_parameter(base_tensor_shape)
-        self.u_graph_neighbor_messages = self._get_parameter(base_tensor_shape)
-        self.linear = to.nn.Linear(self.fully_connected_layer_input_size, self.fully_connected_layer_output_size)
-        self.sigmoid = to.nn.Sigmoid()
-
-    @classmethod
-    def of(cls,
-           time_steps: int,
-           number_of_nodes: int,
-           number_of_node_features: int,
-           fully_connected_layer_input_size: int,
-           fully_connected_layer_output_size: int):
-        return cls(time_steps,
-                   number_of_nodes,
-                   number_of_node_features,
-                   fully_connected_layer_input_size,
-                   fully_connected_layer_output_size)
-
-    def forward(self,
-                node_features: to.Tensor,
-                adjacency_matrix: to.Tensor,
-                batch_size: int) -> to.Tensor:
-        outputs = to.zeros(batch_size, self.fully_connected_layer_output_size)
-        for batch in range(batch_size):
-            outputs[batch] = self.sigmoid(
-                self.linear(
-                    DataPreprocessor.flatten(
-                        self.encode(node_features[batch], adjacency_matrix[batch]),
-                        self.fully_connected_layer_input_size)))
-        return outputs
-
-    def encode(self, node_features: to.Tensor, adjacency_matrix: to.Tensor) -> to.Tensor:
-        messages = self._send_messages(node_features, adjacency_matrix)
-        encodings = self._encode_nodes(node_features, messages)
-        return encodings
-
-    def _send_messages(self, node_features: to.Tensor, adjacency_matrix: to.Tensor) -> to.Tensor:
-        messages = to.zeros((self.number_of_nodes, self.number_of_nodes, self.number_of_node_features))
-        for step in range(self.time_steps):
-            messages = self._compose_messages(node_features, adjacency_matrix, messages)
-        return messages
-
-    def _encode_nodes(self, node_features: to.Tensor, messages: to.Tensor) -> to.Tensor:
-        encoded_node = to.zeros(self.number_of_nodes, self.number_of_node_features)
-        for node_id in range(self.number_of_nodes):
-            encoded_node[node_id] = self._apply_recurrent_layer(node_features, messages, node_id)
-        return encoded_node
-
-    def _apply_recurrent_layer(self, node_features: to.Tensor, messages: to.Tensor, node_id: int) -> to.Tensor:
-        node_encoding_features = self.u_graph_node_features.matmul(node_features[node_id])
-        node_encoding_messages = self.u_graph_neighbor_messages.matmul(to.sum(messages[node_id], dim=0))
-        return to.relu(node_encoding_features + node_encoding_messages)
-
-    def _compose_messages(self,
-                          node_features: to.Tensor,
-                          adjacency_matrix: to.Tensor,
-                          messages: to.Tensor) -> to.Tensor:
-        new_messages = to.zeros(messages.shape)
-        for node_id in range(self.number_of_nodes):
-            node = self._create_node(node_features, adjacency_matrix, node_id)
-            for end_node_id in node.neighbors:
-                message = self._get_message_value(messages, node, end_node_id, node_features)
-                new_messages[node_id, end_node_id] = message
-                if Node.is_adjacency_matrix_symmetric(adjacency_matrix):
-                    new_messages[end_node_id, node_id] = message
-        return new_messages
-
-    def _get_message_value(self,
-                           messages: to.Tensor,
-                           node: Node,
-                           end_node_id: int,
-                           node_features: to.Tensor) -> to.Tensor:
-        previous_messages = self._sum_messages_from_neighbors_except_target(messages,
-                                                                            node,
-                                                                            end_node_id)
-        update_gate = self._pass_through_update_gate(messages, node, end_node_id, node_features)
-        current_memory = self._get_current_memory_message(messages, node, end_node_id, node_features)
-        return to.add(
-            to.mul(
-                to.sub(to.ones(update_gate.shape),
-                       update_gate),
-                previous_messages),
-            to.mul(update_gate,
-                   current_memory))
-
-    def _sum_messages_from_neighbors_except_target(self,
-                                                   messages: to.Tensor,
-                                                   node: Node,
-                                                   end_node_id: int) -> to.Tensor:
-        messages_from_the_other_neighbors = to.zeros(node.features.shape[0])
-        if node.neighbors_count > 1:
-            neighbors_slice = node.get_start_node_neighbors_without_end_node(end_node_id)
-            messages_from_the_other_neighbors = to.sum(messages[neighbors_slice], dim=0)
-        return messages_from_the_other_neighbors
-
-    def _pass_through_update_gate(self,
-                                  messages: to.Tensor,
-                                  node: Node,
-                                  end_node_id: int,
-                                  node_features: to.Tensor) -> to.Tensor:
-        message_from_a_neighbor_other_than_target = self._sum_messages_from_neighbors_except_target(
-            messages,
-            node,
-            end_node_id)
-        update_gate_output = to.sigmoid(
-            to.add(
-                to.add(self.w_gru_update_gate_features.matmul(node_features[node.node_id]),
-                       self.u_gru_update_gate.matmul(
-                           message_from_a_neighbor_other_than_target)),
-                self.b_gru_update_gate))
-        return update_gate_output
-
-    def _get_current_memory_message(self,
-                                    messages: to.Tensor,
-                                    node: Node,
-                                    end_node_id: int,
-                                    node_features: to.Tensor) -> to.Tensor:
-        messages_passed_through_reset_gate = self._keep_or_reset_messages(messages,
-                                                                          node,
-                                                                          end_node_id,
-                                                                          node_features)
-        current_memory_message = to.add(
-            to.add(self.w_gru_current_memory_message_features.matmul(
-                node_features[node.node_id]),
-                self.u_gru_current_memory_message.matmul(
-                    messages_passed_through_reset_gate)),
-            self.b_gru_current_memory_message)
-        return to.tanh(current_memory_message)
-
-    def _keep_or_reset_messages(self,
-                                messages: to.Tensor,
-                                node: Node,
-                                end_node_id: int,
-                                node_features: to.Tensor) -> to.Tensor:
-        return self.u_gru_current_memory_message.matmul(sum([to.mul(to.sigmoid(
-            to.add(
-                to.add(self.w_gru_update_gate_features.matmul(node_features[node.node_id]),
-                       self.u_gru_update_gate.matmul(
-                           messages[node.node_id, reset_node_index])),
-                self.b_gru_update_gate)).long(), messages[node.node_id, reset_node_index])
-                                                             for reset_node_index in
-                                                             node.get_start_node_neighbors_without_end_node(
-                                                                 end_node_id)[0]]))
-
-    def _pass_through_reset_gate(self,
-                                 messages: to.Tensor,
-                                 node: Node,
-                                 end_node_id: int,
-                                 node_features: to.Tensor) -> to.Tensor:
-        return to.sigmoid(
-            to.add(
-                to.add(self.w_gru_update_gate_features.matmul(node_features[node.node_id]),
-                       self.u_gru_update_gate.matmul(messages[node.node_id, end_node_id])),
-                self.b_gru_update_gate)).long()
-
-    @staticmethod
-    def _get_parameter(tensor_shape: List[int]) -> nn.Parameter:
-        return nn.Parameter(nn.init.kaiming_normal_(to.zeros(tensor_shape)), requires_grad=True)
-
-    @staticmethod
-    def _create_node(node_features: to.Tensor, adjacency_matrix: to.Tensor, node_id: int) -> Node:
-        return Node(node_features, adjacency_matrix, node_id)
diff --git a/message_passing_nn/graph/node.py b/message_passing_nn/graph/node.py
deleted file mode 100644
index 3c75408..0000000
--- a/message_passing_nn/graph/node.py
+++ /dev/null
@@ -1,40 +0,0 @@
-from typing import Tuple, List
-
-import torch as to
-
-
-class Node:
-    def __init__(self, all_node_features: to.Tensor, adjacency_matrix: to.Tensor, node_id: int):
-        self.node_id = int(node_id)
-        self.features = self._get_features_of_specific_node(all_node_features)
-        self.symmetry = self.is_adjacency_matrix_symmetric(adjacency_matrix)
-        self.all_neighbors = self._get_all_neighbors(adjacency_matrix)
-        self.neighbors = self._get_neighbors(adjacency_matrix)
-        self.neighbors_count = len(self.all_neighbors)
-
-    @staticmethod
-    def is_adjacency_matrix_symmetric(adjacency_matrix: to.Tensor) -> bool:
-        return to.allclose(adjacency_matrix, adjacency_matrix.t())
-
-    def _get_features_of_specific_node(self, all_node_features: to.Tensor) -> to.Tensor:
-        return all_node_features[self.node_id]
-
-    def _get_all_neighbors(self, adjacency_matrix: to.Tensor) -> to.Tensor:
-        return to.nonzero(adjacency_matrix[self.node_id], as_tuple=True)[0]
-
-    def _get_neighbors(self, adjacency_matrix: to.Tensor) -> List:
-        neighbors = self._get_all_neighbors(adjacency_matrix)
-        if self.symmetry:
-            neighbors = self._get_greater_neighbors(neighbors)
-        return neighbors
-
-    def _get_greater_neighbors(self, neighbors: List) -> List:
-        return [neighbor for neighbor in neighbors if neighbor > self.node_id]
-
-    def get_start_node_neighbors_without_end_node(self, end_node_id: int) -> Tuple:
-        return self._remove_end_node_from_start_node_neighbors(end_node_id), [self.node_id]
-
-    def _remove_end_node_from_start_node_neighbors(self, end_node_id: int) -> to.tensor:
-        end_node_index = (self.all_neighbors == end_node_id).nonzero()[0][0].item()
-        return to.cat((self.all_neighbors[:end_node_index],
-                       self.all_neighbors[end_node_index + 1:])).tolist()
\ No newline at end of file
diff --git a/message_passing_nn/graph/rnn_encoder.cpp b/message_passing_nn/graph/rnn_encoder.cpp
index d08cb1c..14eafab 100644
--- a/message_passing_nn/graph/rnn_encoder.cpp
+++ b/message_passing_nn/graph/rnn_encoder.cpp
@@ -3,80 +3,87 @@
 #include "../utils/derivatives.h"
 #include <iostream>
 
-std::vector<torch::Tensor> forward_cpp(
-    const int& time_steps,
-    const int& number_of_nodes,
-    const int& number_of_node_features,
-    const int& fully_connected_layer_output_size,
-    const int& batch_size,
-    const torch::Tensor& node_features,
-    const torch::Tensor& adjacency_matrix,
-    const torch::Tensor& w_graph_node_features,
-    const torch::Tensor& w_graph_neighbor_messages,
-    const torch::Tensor& u_graph_node_features,
-    const torch::Tensor& u_graph_neighbor_messages,
-    const torch::Tensor& linear_weight,
-    const torch::Tensor& linear_bias) {
+std::vector<at::Tensor> forward_cpp(
+    const at::Tensor& time_steps,
+    const at::Tensor& number_of_nodes,
+    const at::Tensor& number_of_node_features,
+    const at::Tensor& fully_connected_layer_output_size,
+    const at::Tensor& batch_size,
+    const at::Tensor& node_features,
+    const at::Tensor& all_neighbors,
+    const at::Tensor& w_graph_node_features,
+    const at::Tensor& w_graph_neighbor_messages,
+    const at::Tensor& u_graph_node_features,
+    const at::Tensor& u_graph_neighbor_messages,
+    const at::Tensor& linear_weight,
+    const at::Tensor& linear_bias) {
       
-    auto outputs = torch::zeros({batch_size, fully_connected_layer_output_size});
-    auto linear_outputs = torch::zeros({batch_size, fully_connected_layer_output_size});
-    auto messages = torch::zeros({batch_size, number_of_nodes, number_of_nodes, number_of_node_features});
-    auto messages_previous_step = torch::zeros({batch_size, number_of_nodes, number_of_nodes, number_of_node_features});
-    auto node_encoding_messages = torch::zeros({batch_size, number_of_nodes, number_of_node_features});
-    auto encodings = torch::zeros({batch_size, number_of_nodes*number_of_node_features});
+    auto time_steps_int = time_steps.item<int>();
+    auto number_of_nodes_int = number_of_nodes.item<int>();
+    auto number_of_node_features_int = number_of_node_features.item<int>();
+    auto fully_connected_layer_output_size_int = fully_connected_layer_output_size.item<int>();
+    auto batch_size_int = batch_size.item<int>();
+    auto outputs = at::zeros({batch_size_int, fully_connected_layer_output_size_int});
+    auto linear_outputs = at::zeros({batch_size_int, fully_connected_layer_output_size_int});
+    auto messages = at::zeros({batch_size_int, number_of_nodes_int, number_of_nodes_int, number_of_node_features_int});
+    auto messages_previous_step = at::zeros({batch_size_int, number_of_nodes_int, number_of_nodes_int, number_of_node_features_int});
+    auto node_encoding_messages = at::zeros({batch_size_int, number_of_nodes_int, number_of_node_features_int});
+    auto encodings = at::zeros({batch_size_int, number_of_nodes_int*number_of_node_features_int});
+    auto base_messages = at::matmul(w_graph_node_features, node_features);
       
-    for (int batch = 0; batch<batch_size; batch++) {
-      auto messages_vector = compose_messages(time_steps,
-                                        number_of_nodes,
-                                        number_of_node_features,
+    for (int batch = 0; batch<batch_size_int; batch++) {
+      
+      auto messages_vector = compose_messages(time_steps_int,
+                                        number_of_nodes_int,
+                                        number_of_node_features_int,
                                         w_graph_node_features,
                                         w_graph_neighbor_messages,
-                                        node_features[batch],
-                                        adjacency_matrix[batch],
+                                        base_messages[batch],
+                                        all_neighbors[batch],
                                         messages[batch]);
       messages[batch] = messages_vector[0];
       messages_previous_step[batch] = messages_vector[1];
-      encodings[batch] = encode_messages(number_of_nodes,
+      encodings[batch] = encode_messages(number_of_nodes_int,
                                         node_encoding_messages[batch],
                                         u_graph_node_features,
                                         u_graph_neighbor_messages,
                                         node_features[batch],
-                                        adjacency_matrix[batch],
-                                        torch::relu(messages[batch])).view({-1});
-      linear_outputs[batch] = torch::add(linear_bias, torch::matmul(linear_weight, encodings[batch]));
-      outputs[batch] = torch::sigmoid(linear_outputs[batch]);
+                                        all_neighbors[batch],
+                                        messages[batch]).view({-1});
+      linear_outputs[batch] = at::add(at::matmul(linear_weight, encodings[batch]), linear_bias);
+      outputs[batch] = at::sigmoid(linear_outputs[batch]);
     }
     return {outputs, linear_outputs, encodings, messages, messages_previous_step};
   }
 
-std::vector<torch::Tensor> backward_cpp(
-  const torch::Tensor& grad_output,
-  const torch::Tensor& outputs,
-  const torch::Tensor& linear_outputs,
-  const torch::Tensor& encodings,
-  const torch::Tensor& messages_summed,
-  const torch::Tensor& messages_previous_step_summed,
-  const torch::Tensor& messages,
-  const torch::Tensor& node_features,
-  const torch::Tensor& batch_size,
-  const torch::Tensor& number_of_nodes,
-  const torch::Tensor& number_of_node_features,
-  const torch::Tensor& u_graph_neighbor_messages_summed,
-  const torch::Tensor& linear_weight,
-  const torch::Tensor& linear_bias) {
+std::vector<at::Tensor> backward_cpp(
+  const at::Tensor& grad_output,
+  const at::Tensor& outputs,
+  const at::Tensor& linear_outputs,
+  const at::Tensor& encodings,
+  const at::Tensor& messages_summed,
+  const at::Tensor& messages_previous_step_summed,
+  const at::Tensor& messages,
+  const at::Tensor& node_features,
+  const at::Tensor& batch_size,
+  const at::Tensor& number_of_nodes,
+  const at::Tensor& number_of_node_features,
+  const at::Tensor& u_graph_neighbor_messages_summed,
+  const at::Tensor& linear_weight,
+  const at::Tensor& linear_bias) {
   
+
   auto delta_1 = grad_output*d_sigmoid(linear_outputs);
   auto d_linear_bias = delta_1;
-  auto d_linear_weight = torch::matmul(delta_1.transpose(0, 1), encodings);
-  
-  auto delta_2 = torch::matmul(delta_1, linear_weight).reshape({batch_size.item<int>(), number_of_nodes.item<int>(), number_of_node_features.item<int>()})*(d_relu_2d(encodings).reshape({batch_size.item<int>(), number_of_nodes.item<int>(), number_of_node_features.item<int>()}));
-  auto d_u_graph_node_features = torch::matmul(delta_2, node_features.transpose(1, 2));
-  auto d_u_graph_neighbor_messages = torch::matmul(delta_2.transpose(1, 2), messages_summed);
+  auto d_linear_weight = at::matmul(delta_1.transpose(0, 1), encodings);
 
-  auto delta_3 = torch::matmul(delta_2.transpose(1, 2), torch::matmul(u_graph_neighbor_messages_summed, d_relu_4d(messages).transpose(2, 3)));
-  auto d_w_graph_node_features = torch::matmul(delta_3, node_features);
-  auto d_w_graph_neighbor_messages = torch::matmul(delta_3, messages_previous_step_summed);
+  auto delta_2 = at::matmul(delta_1, linear_weight).reshape({batch_size.item<int>(), number_of_nodes.item<int>(), number_of_node_features.item<int>()})*(d_relu_2d(encodings).reshape({batch_size.item<int>(), number_of_nodes.item<int>(), number_of_node_features.item<int>()}));
+  auto d_u_graph_node_features = at::matmul(delta_2, node_features.transpose(1, 2));
+  auto d_u_graph_neighbor_messages = at::matmul(delta_2.transpose(1, 2), messages_summed);
 
+  auto delta_3 = at::matmul(delta_2.transpose(1, 2), at::matmul(u_graph_neighbor_messages_summed, d_relu_4d(messages).transpose(2, 3)));
+  auto d_w_graph_node_features = at::matmul(delta_3.transpose(1, 2), node_features.transpose(1, 2));
+  auto d_w_graph_neighbor_messages = at::matmul(delta_3.transpose(1, 2), messages_previous_step_summed.transpose(1, 2));
 
   return {d_w_graph_node_features, 
           d_w_graph_neighbor_messages, 
diff --git a/message_passing_nn/graph/rnn_encoder.py b/message_passing_nn/graph/rnn_encoder.py
index a2aa081..1940536 100644
--- a/message_passing_nn/graph/rnn_encoder.py
+++ b/message_passing_nn/graph/rnn_encoder.py
@@ -1,6 +1,10 @@
 from typing import Tuple
 
 import rnn_encoder_cpp as rnn_encoder_cpp
+try:
+    import rnn_encoder_cuda_cpp as rnn_encoder_cuda_cpp
+except:
+    pass
 import math
 import torch as to
 import torch.nn as nn
@@ -15,22 +19,27 @@ def forward(ctx,
                 number_of_node_features: int,
                 fully_connected_layer_output_size: int,
                 batch_size: int,
+                device: str,
                 node_features: to.Tensor,
-                adjacency_matrix: to.Tensor,
+                all_neighbors: to.Tensor,
                 w_graph_node_features: to.Tensor,
                 w_graph_neighbor_messages: to.Tensor,
                 u_graph_node_features: to.Tensor,
                 u_graph_neighbor_messages: to.Tensor,
                 linear_weight: to.Tensor,
                 linear_bias: to.Tensor) -> to.Tensor:
-        outputs, linear_outputs, encodings, messages, messages_previous_step = rnn_encoder_cpp.forward(
-            time_steps,
-            number_of_nodes,
-            number_of_node_features,
-            fully_connected_layer_output_size,
-            batch_size,
+        if device == "cuda":
+            cpp_extension = rnn_encoder_cuda_cpp
+        else:
+            cpp_extension = rnn_encoder_cpp
+        outputs, linear_outputs, encodings, messages, messages_previous_step = cpp_extension.forward(
+            to.tensor(time_steps, device=device),
+            to.tensor(number_of_nodes, device=device),
+            to.tensor(number_of_node_features, device=device),
+            to.tensor(fully_connected_layer_output_size, device=device),
+            to.tensor(batch_size, device=device),
             node_features,
-            adjacency_matrix,
+            all_neighbors,
             w_graph_node_features,
             w_graph_neighbor_messages,
             u_graph_node_features,
@@ -39,14 +48,14 @@ def forward(ctx,
             linear_bias)
         variables = [outputs,
                      linear_outputs,
-                     encodings.view(batch_size, number_of_nodes*number_of_node_features),
-                     to.sum(to.relu(messages), dim=2).squeeze(),
-                     to.sum(to.relu(messages_previous_step), dim=2).squeeze(),
+                     encodings.view(batch_size, number_of_nodes * number_of_node_features),
+                     to.sum(to.relu(messages), dim=2),
+                     to.sum(to.relu(messages_previous_step), dim=2),
                      messages,
                      node_features,
-                     to.Tensor([batch_size]),
-                     to.Tensor([number_of_nodes]),
-                     to.Tensor([number_of_node_features]),
+                     to.tensor([batch_size]),
+                     to.tensor([number_of_nodes]),
+                     to.tensor([number_of_node_features]),
                      to.sum(u_graph_neighbor_messages, dim=0),
                      linear_weight,
                      linear_bias]
@@ -61,13 +70,18 @@ def backward(ctx, grad_outputs: to.Tensor) -> Tuple[None,
                                                         None,
                                                         None,
                                                         None,
+                                                        None,
                                                         to.Tensor,
                                                         to.Tensor,
                                                         to.Tensor,
                                                         to.Tensor,
                                                         to.Tensor,
                                                         to.Tensor]:
-        backward_outputs = rnn_encoder_cpp.backward(grad_outputs.contiguous(), *ctx.saved_tensors)
+        if grad_outputs.device == "cuda":
+            cpp_extension = rnn_encoder_cuda_cpp
+        else:
+            cpp_extension = rnn_encoder_cpp
+        backward_outputs = cpp_extension.backward(grad_outputs.contiguous(), *ctx.saved_tensors)
         d_w_graph_node_features, d_w_graph_neighbor_messages, d_u_graph_neighbor_messages, d_u_graph_node_features, d_linear_weight, d_linear_bias = backward_outputs
         return None, \
                None, \
@@ -76,6 +90,7 @@ def backward(ctx, grad_outputs: to.Tensor) -> Tuple[None,
                None, \
                None, \
                None, \
+               None, \
                d_w_graph_node_features, \
                d_w_graph_neighbor_messages, \
                d_u_graph_neighbor_messages, \
@@ -90,7 +105,8 @@ def __init__(self,
                  number_of_nodes: int,
                  number_of_node_features: int,
                  fully_connected_layer_input_size: int,
-                 fully_connected_layer_output_size: int) -> None:
+                 fully_connected_layer_output_size: int,
+                 device: str = "cpu") -> None:
         super(RNNEncoder, self).__init__()
 
         self.time_steps = time_steps
@@ -98,24 +114,31 @@ def __init__(self,
         self.number_of_node_features = number_of_node_features
         self.fully_connected_layer_input_size = fully_connected_layer_input_size
         self.fully_connected_layer_output_size = fully_connected_layer_output_size
+        self.device = device
 
         self.w_graph_node_features = nn.Parameter(
-            to.empty([number_of_node_features, number_of_node_features]),
+            to.empty([number_of_nodes, number_of_nodes],
+                     device=self.device),
             requires_grad=True)
         self.w_graph_neighbor_messages = nn.Parameter(
-            to.empty([number_of_node_features, number_of_node_features]),
+            to.empty([number_of_nodes, number_of_nodes],
+                     device=self.device),
             requires_grad=True)
         self.u_graph_node_features = nn.Parameter(
-            to.empty([number_of_nodes, number_of_nodes]),
+            to.empty([number_of_nodes, number_of_nodes],
+                     device=self.device),
             requires_grad=True)
         self.u_graph_neighbor_messages = nn.Parameter(
-            to.empty([number_of_node_features, number_of_node_features]),
+            to.empty([number_of_node_features, number_of_node_features],
+                     device=self.device),
             requires_grad=True)
         self.linear_weight = nn.Parameter(
-            to.empty([self.fully_connected_layer_output_size, self.fully_connected_layer_input_size]),
+            to.empty([self.fully_connected_layer_output_size, self.fully_connected_layer_input_size],
+                     device=self.device),
             requires_grad=True)
         self.linear_bias = nn.Parameter(
-            to.empty(self.fully_connected_layer_output_size),
+            to.empty(self.fully_connected_layer_output_size,
+                     device=self.device),
             requires_grad=True)
         self.reset_parameters()
 
@@ -130,18 +153,27 @@ def reset_parameters(self):
 
     def forward(self,
                 node_features: to.Tensor,
-                adjacency_matrix: to.Tensor,
+                all_neighbors: to.Tensor,
                 batch_size: int) -> to.Tensor:
         return RNNEncoderFunction.apply(self.time_steps,
                                         self.number_of_nodes,
                                         self.number_of_node_features,
                                         self.fully_connected_layer_output_size,
                                         batch_size,
+                                        self.device,
                                         node_features,
-                                        adjacency_matrix,
+                                        all_neighbors,
                                         self.w_graph_node_features,
                                         self.w_graph_neighbor_messages,
                                         self.u_graph_node_features,
                                         self.u_graph_neighbor_messages,
                                         self.linear_weight,
                                         self.linear_bias)
+
+    def get_model_size(self) -> str:
+        return str(int((self.w_graph_node_features.element_size() * self.w_graph_node_features.nelement() +
+                        self.w_graph_neighbor_messages.element_size() * self.w_graph_neighbor_messages.nelement() +
+                        self.u_graph_node_features.element_size() * self.u_graph_node_features.nelement() +
+                        self.u_graph_neighbor_messages.element_size() * self.u_graph_neighbor_messages.nelement() +
+                        self.linear_weight.element_size() * self.linear_weight.nelement() +
+                        self.linear_bias.element_size() * self.linear_bias.nelement()) * 0.000001))
diff --git a/message_passing_nn/graph/rnn_encoder_cuda.cpp b/message_passing_nn/graph/rnn_encoder_cuda.cpp
new file mode 100644
index 0000000..33be10d
--- /dev/null
+++ b/message_passing_nn/graph/rnn_encoder_cuda.cpp
@@ -0,0 +1,133 @@
+#include <torch/extension.h>
+
+#include <vector>
+
+std::vector<torch::Tensor> forward_cuda_cpp(
+    const at::Tensor& time_steps,
+    const at::Tensor& number_of_nodes,
+    const at::Tensor& number_of_node_features,
+    const at::Tensor& fully_connected_layer_output_size,
+    const at::Tensor& batch_size,
+    const at::Tensor& node_features,
+    const at::Tensor& all_neighbors,
+    const at::Tensor& w_graph_node_features,
+    const at::Tensor& w_graph_neighbor_messages,
+    const at::Tensor& u_graph_node_features,
+    const at::Tensor& u_graph_neighbor_messages,
+    const at::Tensor& linear_weight,
+    const at::Tensor& linear_bias);
+
+std::vector<torch::Tensor> backward_cuda_cpp(
+    const at::Tensor& grad_output,
+    const at::Tensor& outputs,
+    const at::Tensor& linear_outputs,
+    const at::Tensor& encodings,
+    const at::Tensor& messages_summed,
+    const at::Tensor& messages_previous_step_summed,
+    const at::Tensor& messages,
+    const at::Tensor& node_features,
+    const at::Tensor& batch_size,
+    const at::Tensor& number_of_nodes,
+    const at::Tensor& number_of_node_features,
+    const at::Tensor& u_graph_neighbor_messages_summed,
+    const at::Tensor& linear_weight,
+    const at::Tensor& linear_bias);
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> forward_cpp(
+    const at::Tensor& time_steps,
+    const at::Tensor& number_of_nodes,
+    const at::Tensor& number_of_node_features,
+    const at::Tensor& fully_connected_layer_output_size,
+    const at::Tensor& batch_size,
+    const at::Tensor& node_features,
+    const at::Tensor& all_neighbors,
+    const at::Tensor& w_graph_node_features,
+    const at::Tensor& w_graph_neighbor_messages,
+    const at::Tensor& u_graph_node_features,
+    const at::Tensor& u_graph_neighbor_messages,
+    const at::Tensor& linear_weight,
+    const at::Tensor& linear_bias) {
+    CHECK_INPUT(time_steps);
+    CHECK_INPUT(number_of_nodes);
+    CHECK_INPUT(number_of_node_features);
+    CHECK_INPUT(fully_connected_layer_output_size);
+    CHECK_INPUT(batch_size);
+    CHECK_INPUT(node_features);
+    CHECK_INPUT(all_neighbors);
+    CHECK_INPUT(w_graph_node_features);
+    CHECK_INPUT(w_graph_neighbor_messages);
+    CHECK_INPUT(u_graph_node_features);
+    CHECK_INPUT(u_graph_neighbor_messages);
+    CHECK_INPUT(linear_weight);
+    CHECK_INPUT(linear_bias);
+
+  return forward_cuda_cpp(time_steps,
+    number_of_nodes,
+    number_of_node_features,
+    fully_connected_layer_output_size,
+    batch_size,
+    node_features,
+    all_neighbors,
+    w_graph_node_features,
+    w_graph_neighbor_messages,
+    u_graph_node_features,
+    u_graph_neighbor_messages,
+    linear_weight,
+    linear_bias);
+}
+
+std::vector<torch::Tensor> backward_cpp(
+    const at::Tensor& grad_output,
+    const at::Tensor& outputs,
+    const at::Tensor& linear_outputs,
+    const at::Tensor& encodings,
+    const at::Tensor& messages_summed,
+    const at::Tensor& messages_previous_step_summed,
+    const at::Tensor& messages,
+    const at::Tensor& node_features,
+    const at::Tensor& batch_size,
+    const at::Tensor& number_of_nodes,
+    const at::Tensor& number_of_node_features,
+    const at::Tensor& u_graph_neighbor_messages_summed,
+    const at::Tensor& linear_weight,
+    const at::Tensor& linear_bias) {
+    CHECK_INPUT(grad_output);
+    CHECK_INPUT(outputs);
+    CHECK_INPUT(linear_outputs);
+    CHECK_INPUT(encodings);
+    CHECK_INPUT(messages_summed);
+    CHECK_INPUT(messages_previous_step_summed);
+    CHECK_INPUT(messages);
+    CHECK_INPUT(node_features);
+    CHECK_INPUT(batch_size);
+    CHECK_INPUT(number_of_nodes);
+    CHECK_INPUT(number_of_node_features);
+    CHECK_INPUT(u_graph_neighbor_messages_summed);
+    CHECK_INPUT(linear_weight);
+    CHECK_INPUT(linear_bias);
+
+  return backward_cuda_cpp(
+      grad_output,
+      outputs,
+      linear_outputs,
+      encodings,
+      messages_summed,
+      messages_previous_step_summed,
+      messages,
+      node_features,
+      batch_size,
+      number_of_nodes,
+      number_of_node_features,
+      u_graph_neighbor_messages_summed,
+      linear_weight,
+      linear_bias);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &forward_cpp, "RNN encoder forward pass (CUDA)");
+  m.def("backward", &backward_cpp, "RNN encoder backward pass (CUDA)");
+}
\ No newline at end of file
diff --git a/message_passing_nn/graph/rnn_encoder_cuda_kernel.cu b/message_passing_nn/graph/rnn_encoder_cuda_kernel.cu
new file mode 100644
index 0000000..9e8e794
--- /dev/null
+++ b/message_passing_nn/graph/rnn_encoder_cuda_kernel.cu
@@ -0,0 +1,158 @@
+#include <torch/extension.h>
+#include "../utils/derivatives.h"
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <vector>
+
+
+template <typename scalar_t>
+__global__ void compose_messages_kernel(
+    torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> base_neighbor_messages,
+    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> w_graph_neighbor_messages,
+    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> base_messages,
+    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> all_neighbors,
+    torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> new_messages) {
+
+    const int index = threadIdx.x;
+    const int stride = blockDim.x;
+
+    for (int node_id = index; node_id < all_neighbors.size(0); node_id += stride) {
+      for (int end_node_index = 0; end_node_index < all_neighbors.size(1); end_node_index++){
+        auto end_node_id = std::round(all_neighbors[node_id][end_node_index]);
+        if (end_node_id >= 0) {
+          for (int index_feature = 0; index_feature < new_messages.size(2); index_feature++) {
+            new_messages[node_id][end_node_id][index_feature] += base_messages[node_id][index_feature];
+          }
+          for (int neighbor_index = 0; neighbor_index < all_neighbors.size(1); neighbor_index++) {
+            auto neighbor = std::round(all_neighbors[node_id][neighbor_index]);
+            if (neighbor >= 0 && neighbor_index!=end_node_index) {
+              for (int index_feature = 0; index_feature < new_messages.size(2); index_feature++) {
+                new_messages[node_id][end_node_id][index_feature] += base_neighbor_messages[neighbor][node_id][index_feature];
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
+at::Tensor encode_messages(
+    const int& number_of_nodes,
+    const at::Tensor& node_encoding_messages,
+    const at::Tensor& u_graph_node_features,
+    const at::Tensor& u_graph_neighbor_messages,
+    const at::Tensor& node_features,
+    const at::Tensor& all_neighbors,
+    const at::Tensor& messages) {
+
+    for (int node_id = 0; node_id<number_of_nodes; node_id++) {
+      for (int end_node_index = 0; end_node_index<all_neighbors.sizes()[1]; end_node_index++){
+        auto end_node_id = all_neighbors[node_id][end_node_index].item<int64_t>();
+        if (end_node_id >= 0) {
+          node_encoding_messages[node_id] += at::matmul(u_graph_neighbor_messages, at::relu(messages[end_node_id][node_id]));
+        }
+      }
+    }
+    return at::relu(at::add(at::matmul(u_graph_node_features, node_features), node_encoding_messages));
+  }
+
+std::vector<at::Tensor> forward_cuda_cpp(
+    const at::Tensor& time_steps,
+    const at::Tensor& number_of_nodes,
+    const at::Tensor& number_of_node_features,
+    const at::Tensor& fully_connected_layer_output_size,
+    const at::Tensor& batch_size,
+    const at::Tensor& node_features,
+    const at::Tensor& all_neighbors,
+    const at::Tensor& w_graph_node_features,
+    const at::Tensor& w_graph_neighbor_messages,
+    const at::Tensor& u_graph_node_features,
+    const at::Tensor& u_graph_neighbor_messages,
+    const at::Tensor& linear_weight,
+    const at::Tensor& linear_bias) {
+      
+    auto outputs = at::zeros({batch_size.item<int>(), fully_connected_layer_output_size.item<int>()}, at::kCUDA);
+    auto linear_outputs = at::zeros({batch_size.item<int>(), fully_connected_layer_output_size.item<int>()}, at::kCUDA);
+    auto messages = at::zeros({batch_size.item<int>(), number_of_nodes.item<int>(), number_of_nodes.item<int>(), number_of_node_features.item<int>()}, at::kCUDA);
+    auto messages_previous_step = at::zeros({batch_size.item<int>(), number_of_nodes.item<int>(), number_of_nodes.item<int>(), number_of_node_features.item<int>()}, at::kCUDA);
+    auto node_encoding_messages = at::zeros({batch_size.item<int>(), number_of_nodes.item<int>(), number_of_node_features.item<int>()}, at::kCUDA);
+    auto encodings = at::zeros({batch_size.item<int>(), number_of_nodes.item<int>()*number_of_node_features.item<int>()}, at::kCUDA);
+    
+    const int threads = 1024;
+    const dim3 blocks(std::floor(number_of_nodes.item<int>()/threads) + 1);
+
+    auto base_messages = at::matmul(w_graph_node_features, node_features);
+      
+    for (int batch = 0; batch<batch_size.item<int>(); batch++) {
+      auto new_messages = at::zeros_like({messages[batch]}, at::kCUDA);
+      auto previous_messages = at::zeros_like({messages[batch]}, at::kCUDA);
+      const auto number_of_nodes = all_neighbors[batch].size(0);
+      const auto max_neighbors = all_neighbors[batch].size(1);
+      
+      for (int time_step = 0; time_step<time_steps.item<int>(); time_step++) {
+        auto base_neighbor_messages = at::matmul(w_graph_neighbor_messages, at::relu(new_messages));
+        std::swap(previous_messages, new_messages);
+        auto base_messages_of_batch = base_messages[batch];
+        auto neighbors_of_batch = all_neighbors[batch];
+        AT_DISPATCH_FLOATING_TYPES(new_messages.type(), "forward_cpp_cuda", ([&] {
+          compose_messages_kernel<scalar_t><<<blocks, threads>>>(base_neighbor_messages.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>(),
+                                                                 w_graph_neighbor_messages.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+                                                                 base_messages_of_batch.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+                                                                 neighbors_of_batch.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+                                                                 new_messages.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>());
+        }));
+      }
+
+      messages[batch] = new_messages;
+      messages_previous_step[batch] = previous_messages;
+      encodings[batch] = encode_messages(number_of_nodes,
+                                        node_encoding_messages[batch],
+                                        u_graph_node_features,
+                                        u_graph_neighbor_messages,
+                                        node_features[batch],
+                                        all_neighbors[batch],
+                                        messages[batch]).view({-1});
+      linear_outputs[batch] = at::add(at::matmul(linear_weight, encodings[batch]), linear_bias);
+      outputs[batch] = at::sigmoid(linear_outputs[batch]);
+    }
+    return {outputs, linear_outputs, encodings, messages, messages_previous_step};
+  }
+
+std::vector<at::Tensor> backward_cuda_cpp(
+  const at::Tensor& grad_output,
+  const at::Tensor& outputs,
+  const at::Tensor& linear_outputs,
+  const at::Tensor& encodings,
+  const at::Tensor& messages_summed,
+  const at::Tensor& messages_previous_step_summed,
+  const at::Tensor& messages,
+  const at::Tensor& node_features,
+  const at::Tensor& batch_size,
+  const at::Tensor& number_of_nodes,
+  const at::Tensor& number_of_node_features,
+  const at::Tensor& u_graph_neighbor_messages_summed,
+  const at::Tensor& linear_weight,
+  const at::Tensor& linear_bias) {
+  
+  auto delta_1 = grad_output*d_sigmoid(linear_outputs);
+  auto d_linear_bias = delta_1;
+  auto d_linear_weight = at::matmul(delta_1.transpose(0, 1), encodings);
+  
+  auto delta_2 = at::matmul(delta_1, linear_weight).reshape({batch_size.item<int>(), number_of_nodes.item<int>(), number_of_node_features.item<int>()})*(d_relu_2d(encodings).reshape({batch_size.item<int>(), number_of_nodes.item<int>(), number_of_node_features.item<int>()}));
+  auto d_u_graph_node_features = at::matmul(delta_2, node_features.transpose(1, 2));
+  auto d_u_graph_neighbor_messages = at::matmul(delta_2.transpose(1, 2), messages_summed);
+
+  auto delta_3 = at::matmul(delta_2.transpose(1, 2), at::matmul(u_graph_neighbor_messages_summed, d_relu_4d(messages).transpose(2, 3)));
+  auto d_w_graph_node_features = at::matmul(delta_3.transpose(1, 2), node_features.transpose(1, 2));
+  auto d_w_graph_neighbor_messages = at::matmul(delta_3.transpose(1, 2), messages_previous_step_summed.transpose(1, 2));
+
+
+  return {d_w_graph_node_features, 
+          d_w_graph_neighbor_messages, 
+          d_u_graph_node_features, 
+          d_u_graph_neighbor_messages,
+          d_linear_weight,
+          d_linear_bias};
+}
\ No newline at end of file
diff --git a/message_passing_nn/infrastructure/__init__.py b/message_passing_nn/infrastructure/__init__.py
new file mode 100644
index 0000000..090295f
--- /dev/null
+++ b/message_passing_nn/infrastructure/__init__.py
@@ -0,0 +1,1 @@
+from message_passing_nn.infrastructure.file_system_repository import FileSystemRepository
diff --git a/message_passing_nn/infrastructure/file_system_repository.py b/message_passing_nn/infrastructure/file_system_repository.py
new file mode 100644
index 0000000..660af35
--- /dev/null
+++ b/message_passing_nn/infrastructure/file_system_repository.py
@@ -0,0 +1,19 @@
+import logging
+import pickle
+
+import torch as to
+
+
+class FileSystemRepository:
+    def __init__(self, data_directory: str, dataset: str) -> None:
+        super().__init__()
+        self.data_directory = data_directory + dataset + '/'
+        self.test_mode = False
+
+    def save(self, filename: str, data_to_save: to.Tensor) -> None:
+        with open(self.data_directory + filename, 'wb') as file:
+            pickle.dump(data_to_save, file)
+
+    @staticmethod
+    def get_logger() -> logging.Logger:
+        return logging.getLogger('message_passing_nn')
diff --git a/message_passing_nn/infrastructure/graph_dataset.py b/message_passing_nn/infrastructure/graph_dataset.py
new file mode 100644
index 0000000..73ecb52
--- /dev/null
+++ b/message_passing_nn/infrastructure/graph_dataset.py
@@ -0,0 +1,103 @@
+import logging
+import os
+import pickle
+from typing import List, Tuple
+
+import torch as to
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+
+class GraphDataset(Dataset):
+    def __init__(self,
+                 data_directory: str,
+                 test_mode: bool = False) -> None:
+        self.data_directory = data_directory
+        self.test_mode = test_mode
+        self.dataset = self._load_data() if self.data_directory else []
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, index: int) -> Tuple[to.Tensor, to.Tensor, to.Tensor, str]:
+        return self.dataset[index][0], self.dataset[index][1], self.dataset[index][2], self.dataset[index][3]
+
+    def _load_data(self) -> List[Tuple[to.Tensor, to.Tensor, to.Tensor, str]]:
+        self.get_logger().info("Loading dataset")
+        files_in_path = self._extract_name_prefixes_from_filenames()
+        dataset = []
+        size = 0
+        disable_progress_bar = self.test_mode
+        for filename_index in tqdm(range(len(files_in_path)), disable=disable_progress_bar):
+            filename = files_in_path[filename_index]
+            try:
+                dataset.append(
+                    (self._get_features(filename), self._get_all_neighbors(filename), self._get_labels(filename),
+                     filename))
+            except:
+                self.get_logger().info("Skipped " + filename)
+            size += self._get_size(dataset[-1])
+        self.get_logger().info(
+            "Loaded " + str(len(dataset)) + " files. Size: " + str(int(size * 0.000001)) + " MB")
+        return dataset
+
+    @staticmethod
+    def _to_list(dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]]) -> List[Tuple[to.Tensor, to.Tensor]]:
+        return [(dataset[index][0], dataset[index][1]) for index in range(len(dataset))]
+
+    @staticmethod
+    def _extract_labels(dataset: List[Tuple[to.Tensor, to.Tensor, to.Tensor]]) -> List[to.Tensor]:
+        return [dataset[index][2] for index in range(len(dataset))]
+
+    def _get_labels(self, filename: str) -> to.Tensor:
+        with open(self.data_directory + filename + 'labels.pickle', 'rb') as labels_file:
+            labels = pickle.load(labels_file).float()
+        return labels
+
+    def _get_features(self, filename: str) -> to.Tensor:
+        with open(self.data_directory + filename + 'features.pickle', 'rb') as features_file:
+            features = pickle.load(features_file).float()
+        return features
+
+    def _get_all_neighbors(self, filename: str) -> to.Tensor:
+        with open(self.data_directory + filename + 'adjacency-matrix.pickle', 'rb') as adjacency_matrix_file:
+            adjacency_matrix = pickle.load(adjacency_matrix_file).float()
+            number_of_nodes = adjacency_matrix.shape[0]
+            all_neighbors_list = []
+            max_number_of_neighbors = -1
+            for node_id in range(adjacency_matrix.shape[0]):
+                neighbors = to.nonzero(adjacency_matrix[node_id], as_tuple=True)[0].tolist()
+                if len(neighbors) > max_number_of_neighbors:
+                    max_number_of_neighbors = len(neighbors)
+                all_neighbors_list.append(neighbors)
+            all_neighbors = self._get_minus_ones_tensor(number_of_nodes, max_number_of_neighbors)
+            for node_id in range(number_of_nodes):
+                all_neighbors[node_id, :len(all_neighbors_list[node_id])] = to.tensor(all_neighbors_list[node_id])
+        return all_neighbors
+
+    @staticmethod
+    def _get_minus_ones_tensor(number_of_nodes: int, max_number_of_neighbors: int) -> to.Tensor:
+        return to.zeros(number_of_nodes, max_number_of_neighbors) - to.ones(number_of_nodes, max_number_of_neighbors)
+
+    @staticmethod
+    def _get_size(data: Tuple[to.Tensor, to.Tensor, to.Tensor, str]) -> int:
+        return int(data[0].element_size() * data[0].nelement() +
+                   data[1].element_size() * data[1].nelement() +
+                   data[2].element_size() * data[2].nelement())
+
+    def _extract_name_prefixes_from_filenames(self) -> List[str]:
+        return list(set([self._reconstruct_filename(file) for file in self._get_data_filenames()]))
+
+    def _get_data_filenames(self) -> List[str]:
+        return sorted([file for file in os.listdir(self.data_directory) if file.endswith(".pickle")])
+
+    @staticmethod
+    def _reconstruct_filename(file: str) -> str:
+        return "_".join(file.split("_")[:-1]) + "_"
+
+    def enable_test_mode(self) -> None:
+        self.test_mode = True
+
+    @staticmethod
+    def get_logger() -> logging.Logger:
+        return logging.getLogger('message_passing_nn')
diff --git a/message_passing_nn/model/inferencer.py b/message_passing_nn/model/inferencer.py
index e831169..2809dca 100644
--- a/message_passing_nn/model/inferencer.py
+++ b/message_passing_nn/model/inferencer.py
@@ -13,17 +13,16 @@ def __init__(self, data_preprocessor: DataPreprocessor, device: str, normalize:
         self.device = device
         self.normalize = normalize
 
-    def do_inference(self, model: nn.Module, inference_data: DataLoader) -> List[Tuple[to.Tensor, to.Tensor]]:
+    def do_inference(self, model: nn.Module, inference_data: DataLoader) -> List[Tuple[to.Tensor, to.Tensor, str]]:
         outputs_labels_pairs = []
         with to.no_grad():
-            for features, labels in inference_data:
-                node_features, adjacency_matrix = features
-                node_features, adjacency_matrix, labels = node_features.to(self.device), \
-                                                          adjacency_matrix.to(self.device), \
-                                                          labels.to(self.device)
+            for node_features, all_neighbors, labels, tag in inference_data:
+                node_features, all_neighbors, labels = (node_features.to(self.device),
+                                                        all_neighbors.to(self.device),
+                                                        labels.to(self.device))
                 if self.normalize:
                     node_features = self.preprocessor.normalize(node_features, self.device)
                     labels = self.preprocessor.normalize(labels, self.device)
-                outputs = model.forward(node_features, adjacency_matrix, batch_size=1)
-                outputs_labels_pairs.append((outputs, labels))
+                outputs = model.forward(node_features, all_neighbors, batch_size=1)
+                outputs_labels_pairs.append((outputs, labels, tag))
         return outputs_labels_pairs
diff --git a/message_passing_nn/model/loader.py b/message_passing_nn/model/loader.py
index 094c4f6..1716479 100644
--- a/message_passing_nn/model/loader.py
+++ b/message_passing_nn/model/loader.py
@@ -12,8 +12,8 @@ def __init__(self, model: str) -> None:
 
     def load_model(self, data_dimensions: Tuple, path_to_model: str) -> nn.Module:
         model_parameters = self._get_model_parameters_from_path(path_to_model)
-        node_features_size, adjacency_matrix_size, labels_size = data_dimensions
-        number_of_nodes = adjacency_matrix_size[0]
+        node_features_size, labels_size = data_dimensions
+        number_of_nodes = node_features_size[0]
         number_of_node_features = node_features_size[1]
         fully_connected_layer_output_size = labels_size[0]
         self.model = self.model(time_steps=int(model_parameters['time_steps']),
diff --git a/message_passing_nn/model/trainer.py b/message_passing_nn/model/trainer.py
index 1c36e00..dbb0f81 100644
--- a/message_passing_nn/model/trainer.py
+++ b/message_passing_nn/model/trainer.py
@@ -1,8 +1,7 @@
 import logging
-from multiprocessing import get_context
 from typing import Dict, Any, Tuple
-import numpy as np
 
+import numpy as np
 import torch as to
 from torch.nn.modules.module import Module
 from torch.optim.optimizer import Optimizer
@@ -15,11 +14,9 @@
 
 
 class Trainer:
-    def __init__(self, preprocessor: Preprocessor, device: str, cpu_multiprocessing: bool = False,
-                 normalize: bool = False) -> None:
+    def __init__(self, preprocessor: Preprocessor, device: str, normalize: bool = False) -> None:
         self.preprocessor = preprocessor
         self.device = device
-        self.cpu_multiprocessing = cpu_multiprocessing
         self.normalize = normalize
         self.model = None
         self.loss_function = None
@@ -28,8 +25,8 @@ def __init__(self, preprocessor: Preprocessor, device: str, cpu_multiprocessing:
     def instantiate_attributes(self,
                                data_dimensions: Tuple,
                                configuration_dictionary: Dict) -> None:
-        node_features_size, adjacency_matrix_size, labels_size = data_dimensions
-        number_of_nodes = adjacency_matrix_size[0]
+        node_features_size, labels_size = data_dimensions
+        number_of_nodes = node_features_size[0]
         number_of_node_features = node_features_size[1]
         fully_connected_layer_output_size = labels_size[0]
         self.model = ModelSelector.load_model(configuration_dictionary['model'])
@@ -37,35 +34,34 @@ def instantiate_attributes(self,
                                 number_of_nodes=number_of_nodes,
                                 number_of_node_features=number_of_node_features,
                                 fully_connected_layer_input_size=number_of_nodes * number_of_node_features,
-                                fully_connected_layer_output_size=fully_connected_layer_output_size)
+                                fully_connected_layer_output_size=fully_connected_layer_output_size,
+                                device=self.device)
+        self.get_logger().info('Loaded the ' + configuration_dictionary['model'] +
+                               ' model. Model weights size: ' + self.model.get_model_size() + ' MB')
         self.model.to(self.device)
         self.loss_function = self._instantiate_the_loss_function(
             LossFunctionSelector.load_loss_function(configuration_dictionary['loss_function']))
+        self.get_logger().info('Loss function: ' + configuration_dictionary['loss_function'])
         self.optimizer = self._instantiate_the_optimizer(
             OptimizerSelector.load_optimizer(configuration_dictionary['optimizer']))
+        self.get_logger().info('Optimizer: ' + configuration_dictionary['optimizer'])
 
     def do_train(self, training_data: DataLoader, epoch: int) -> float:
-        training_loss = 0.0
-        if self.cpu_multiprocessing and self.device == 'cpu':
-            with get_context("spawn").Pool() as pool:
-                training_loss += np.average(pool.map(self._do_train_batch, training_data))
-        else:
-            training_loss += np.average(list(map(self._do_train_batch, training_data)))
+        training_loss = np.average(list(map(self._do_train_batch, training_data)))
         self.get_logger().info('[Iteration %d] training loss: %.6f' % (epoch, training_loss))
         return training_loss
 
     def _do_train_batch(self, training_data: DataLoader) -> float:
-        features, labels = training_data
-        node_features, adjacency_matrix = features
-        node_features, adjacency_matrix, labels = node_features.to(self.device), \
-                                                  adjacency_matrix.to(self.device), \
-                                                  labels.to(self.device)
+        node_features, all_neighbors, labels, _ = training_data
+        node_features, all_neighbors, labels = (node_features.to(self.device),
+                                                all_neighbors.to(self.device),
+                                                labels.to(self.device))
         current_batch_size = self._get_current_batch_size(labels)
         if self.normalize:
             node_features = self.preprocessor.normalize(node_features, self.device)
             labels = self.preprocessor.normalize(labels, self.device)
         self.optimizer.zero_grad()
-        outputs = self.model.forward(node_features, adjacency_matrix, batch_size=current_batch_size)
+        outputs = self.model(node_features, all_neighbors, batch_size=current_batch_size)
         loss = self.loss_function(outputs, labels)
         self._do_backpropagate(loss)
         return loss.item()
@@ -74,16 +70,15 @@ def do_evaluate(self, evaluation_data: DataLoader, epoch: int = None) -> float:
         with to.no_grad():
             evaluation_loss = []
             if len(evaluation_data):
-                for features_validation, labels_validation in evaluation_data:
-                    node_features, adjacency_matrix = features_validation
-                    node_features, adjacency_matrix, labels_validation = node_features.to(self.device), \
-                                                                         adjacency_matrix.to(self.device), \
-                                                                         labels_validation.to(self.device)
+                for node_features, all_neighbors, labels_validation, _ in evaluation_data:
+                    node_features, all_neighbors, labels_validation = (node_features.to(self.device),
+                                                                       all_neighbors.to(self.device),
+                                                                       labels_validation.to(self.device))
                     if self.normalize:
                         node_features = self.preprocessor.normalize(node_features, self.device)
                         labels_validation = self.preprocessor.normalize(labels_validation, self.device)
                     current_batch_size = self._get_current_batch_size(labels_validation)
-                    outputs = self.model.forward(node_features, adjacency_matrix, current_batch_size)
+                    outputs = self.model(node_features, all_neighbors, current_batch_size)
                     loss = self.loss_function(outputs, labels_validation)
                     evaluation_loss.append(float(loss))
                 evaluation_loss = np.average(evaluation_loss)
diff --git a/message_passing_nn/repository/__init__.py b/message_passing_nn/repository/__init__.py
deleted file mode 100644
index 1b167c8..0000000
--- a/message_passing_nn/repository/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-from message_passing_nn.repository.file_system_repository import FileSystemRepository
-from message_passing_nn.repository.repository import Repository
diff --git a/message_passing_nn/repository/file_system_repository.py b/message_passing_nn/repository/file_system_repository.py
deleted file mode 100644
index 686776e..0000000
--- a/message_passing_nn/repository/file_system_repository.py
+++ /dev/null
@@ -1,56 +0,0 @@
-import logging
-import os
-import pickle
-from typing import List, Tuple
-
-import torch as to
-
-from message_passing_nn.repository.repository import Repository
-
-
-class FileSystemRepository(Repository):
-    def __init__(self, data_directory: str, dataset: str) -> None:
-        super().__init__()
-        self.data_directory = data_directory + dataset + '/'
-
-    def save(self, filename: str, data_to_save: to.Tensor) -> None:
-        with open(self.data_directory + filename, 'wb') as file:
-            pickle.dump(data_to_save, file)
-
-    def get_all_data(self) -> List[Tuple[to.Tensor, to.Tensor, to.Tensor]]:
-        self.get_logger().info("Loading dataset")
-        files_in_path = self._extract_name_prefixes_from_filenames()
-        dataset = []
-        for filename in files_in_path:
-            dataset.append((self._get_features(filename), self._get_adjacency_matrix(filename), self._get_labels(filename)))
-        self.get_logger().info("Loaded " + str(len(dataset)) + " files")
-        return dataset
-
-    def _get_labels(self, filename: str) -> to.Tensor:
-        with open(self.data_directory + filename + 'labels.pickle', 'rb') as labels_file:
-            labels = pickle.load(labels_file).float()
-        return labels
-
-    def _get_features(self, filename: str) -> to.Tensor:
-        with open(self.data_directory + filename + 'features.pickle', 'rb') as features_file:
-            features = pickle.load(features_file).float()
-        return features
-
-    def _get_adjacency_matrix(self, filename: str) -> to.Tensor:
-        with open(self.data_directory + filename + 'adjacency-matrix.pickle', 'rb') as adjacency_matrix_file:
-            features = pickle.load(adjacency_matrix_file).float()
-        return features
-
-    def _extract_name_prefixes_from_filenames(self):
-        return set([self._reconstruct_filename(file) for file in self._get_data_filenames()])
-
-    def _get_data_filenames(self):
-        return sorted([file for file in os.listdir(self.data_directory) if file.endswith(".pickle")])
-
-    @staticmethod
-    def _reconstruct_filename(file):
-        return "_".join(file.split("_")[:-1]) + "_"
-
-    @staticmethod
-    def get_logger() -> logging.Logger:
-        return logging.getLogger('message_passing_nn')
diff --git a/message_passing_nn/repository/repository.py b/message_passing_nn/repository/repository.py
deleted file mode 100644
index 14e30a7..0000000
--- a/message_passing_nn/repository/repository.py
+++ /dev/null
@@ -1,15 +0,0 @@
-from abc import ABCMeta, abstractmethod
-from typing import List, Tuple, Any
-
-
-class Repository(metaclass=ABCMeta):
-    def __init__(self):
-        pass
-
-    @abstractmethod
-    def save(self, filename: str, dataset_name: str) -> None:
-        pass
-
-    @abstractmethod
-    def get_all_data(self) -> List[Tuple[Any, Any, Any]]:
-        pass
diff --git a/message_passing_nn/usecase/grid_search.py b/message_passing_nn/usecase/grid_search.py
index e83d71c..1c7cad0 100644
--- a/message_passing_nn/usecase/grid_search.py
+++ b/message_passing_nn/usecase/grid_search.py
@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import Dict, List, Tuple
 
 import itertools
@@ -6,64 +7,57 @@
 from torch.utils.data.dataloader import DataLoader
 
 from message_passing_nn.data.data_preprocessor import DataPreprocessor
-from message_passing_nn.repository.repository import Repository
+from message_passing_nn.infrastructure.graph_dataset import GraphDataset
 from message_passing_nn.model.trainer import Trainer
-from message_passing_nn.utils.saver import Saver
 from message_passing_nn.usecase import Usecase
+from message_passing_nn.utils.saver import Saver
 
 
 class GridSearch(Usecase):
     def __init__(self,
-                 training_data_repository: Repository,
+                 data_path: str,
                  data_preprocessor: DataPreprocessor,
                  trainer: Trainer,
                  grid_search_dictionary: Dict,
-                 saver: Saver) -> None:
-        self.repository = training_data_repository
+                 saver: Saver,
+                 test_mode: bool = False) -> None:
+        self.data_path = data_path
         self.data_preprocessor = data_preprocessor
         self.trainer = trainer
         self.grid_search_dictionary = grid_search_dictionary
         self.saver = saver
+        self.test_mode = test_mode
 
     def start(self) -> Dict:
         all_grid_search_configurations = self._get_all_grid_search_configurations()
-        self.get_logger().info('Started Training')
         losses = {'training_loss': {},
                   'validation_loss': {},
                   'test_loss': {}}
         configuration_id = ''
         for configuration in all_grid_search_configurations:
-            configuration_id, grid_search_configuration_dictionary = self._get_grid_search_configuration_dictionary(
-                configuration)
-            losses = self._search_configuration(configuration_id,
-                                                grid_search_configuration_dictionary,
-                                                losses)
+            configuration_id, configuration_dictionary = self._get_configuration_dictionary(configuration)
+            losses = self._search_configuration(configuration_id, configuration_dictionary, losses)
         self.saver.save_results(configuration_id, losses)
         self.get_logger().info('Finished Training')
         return losses
 
-    def _search_configuration(self,
-                              configuration_id: str,
-                              grid_search_configuration_dictionary: Dict,
-                              losses: Dict) -> Dict:
-        training_data, validation_data, test_data, data_dimensions = self._prepare_dataset(
-            grid_search_configuration_dictionary)
-        self.trainer.instantiate_attributes(data_dimensions, grid_search_configuration_dictionary)
-        losses = self._update_losses_with_configuration_id(grid_search_configuration_dictionary, losses)
+    def _search_configuration(self, configuration_id: str, configuration_dictionary: Dict, losses: Dict) -> Dict:
+        training_data, validation_data, test_data, data_dimensions = self._prepare_dataset(configuration_dictionary)
+        self.trainer.instantiate_attributes(data_dimensions, configuration_dictionary)
+        losses = self._update_losses_with_configuration_id(configuration_dictionary, losses)
         validation_loss_max = np.inf
-        for epoch in range(1, grid_search_configuration_dictionary['epochs'] + 1):
+        self.get_logger().info('Started Training')
+        for epoch in range(1, configuration_dictionary['epochs'] + 1):
             training_loss = self.trainer.do_train(training_data, epoch)
-            losses['training_loss'][grid_search_configuration_dictionary["configuration_id"]].update(
-                {epoch: training_loss})
-            if epoch % grid_search_configuration_dictionary["validation_period"] == 0:
+            losses['training_loss'][configuration_dictionary["configuration_id"]].update({epoch: training_loss})
+            if epoch % configuration_dictionary["validation_period"] == 0:
                 validation_loss = self.trainer.do_evaluate(validation_data, epoch)
-                losses['validation_loss'][grid_search_configuration_dictionary["configuration_id"]].update(
+                losses['validation_loss'][configuration_dictionary["configuration_id"]].update(
                     {epoch: validation_loss})
                 if validation_loss < validation_loss_max:
                     self.saver.save_model(epoch, configuration_id, self.trainer.model)
         test_loss = self.trainer.do_evaluate(test_data)
-        losses['test_loss'][grid_search_configuration_dictionary["configuration_id"]].update(
-            {"final_epoch": test_loss})
+        losses['test_loss'][configuration_dictionary["configuration_id"]].update({"final_epoch": test_loss})
         return losses
 
     @staticmethod
@@ -74,25 +68,24 @@ def _update_losses_with_configuration_id(configuration_dictionary: Dict, losses:
         return losses
 
     @staticmethod
-    def _get_grid_search_configuration_dictionary(configuration: Tuple[Tuple]) -> Tuple[str, Dict]:
-        grid_search_configuration_dictionary = dict(((key, value) for key, value in configuration))
+    def _get_configuration_dictionary(configuration: Tuple[Tuple]) -> Tuple[str, Dict]:
+        configuration_dictionary = dict(((key, value) for key, value in configuration))
         configuration_id = 'configuration&id'
-        for key, value in grid_search_configuration_dictionary.items():
+        for key, value in configuration_dictionary.items():
             configuration_id += "__" + "&".join([key, str(value)])
-        grid_search_configuration_dictionary.update({"configuration_id": configuration_id})
-        return configuration_id, grid_search_configuration_dictionary
+        configuration_dictionary.update({"configuration_id": configuration_id})
+        return configuration_id, configuration_dictionary
 
     def _prepare_dataset(self, configuration_dictionary: Dict) -> Tuple[DataLoader, DataLoader, DataLoader, Tuple]:
-        raw_dataset = self.repository.get_all_data()
-        equalized_dataset = self.data_preprocessor.equalize_dataset_dimensions(raw_dataset,
-                                                                               configuration_dictionary['maximum_number_of_nodes'],
-                                                                               configuration_dictionary['maximum_number_of_features'])
+        dataset = GraphDataset(self.data_path, test_mode=self.test_mode)
+        dataset.enable_test_mode()
+        self.get_logger().info("Calculating all neighbors for each node")
         training_data, validation_data, test_data = self.data_preprocessor \
-            .train_validation_test_split(equalized_dataset,
+            .train_validation_test_split(dataset,
                                          configuration_dictionary['batch_size'],
                                          configuration_dictionary['validation_split'],
                                          configuration_dictionary['test_split'])
-        data_dimensions = self.data_preprocessor.extract_data_dimensions(equalized_dataset)
+        data_dimensions = self.data_preprocessor.extract_data_dimensions(dataset)
         return training_data, validation_data, test_data, data_dimensions
 
     def _get_all_grid_search_configurations(self) -> List[Tuple[Tuple]]:
diff --git a/message_passing_nn/usecase/inference.py b/message_passing_nn/usecase/inference.py
index 396c283..a1d7760 100644
--- a/message_passing_nn/usecase/inference.py
+++ b/message_passing_nn/usecase/inference.py
@@ -1,28 +1,30 @@
 import logging
+from typing import Tuple
 
 from torch.utils.data.dataloader import DataLoader
-from typing import Tuple
 
 from message_passing_nn.data import DataPreprocessor
+from message_passing_nn.infrastructure.graph_dataset import GraphDataset
 from message_passing_nn.model.inferencer import Inferencer
 from message_passing_nn.model.loader import Loader
-from message_passing_nn.repository import Repository
 from message_passing_nn.usecase import Usecase
 from message_passing_nn.utils import Saver
 
 
 class Inference(Usecase):
     def __init__(self,
-                 training_data_repository: Repository,
+                 data_path: str,
                  data_preprocessor: DataPreprocessor,
                  loader: Loader,
                  inferencer: Inferencer,
-                 saver: Saver) -> None:
-        self.repository = training_data_repository
+                 saver: Saver,
+                 test_mode: bool = False) -> None:
+        self.data_path = data_path
         self.data_preprocessor = data_preprocessor
         self.loader = loader
         self.inferencer = inferencer
         self.saver = saver
+        self.test_mode = test_mode
 
     def start(self) -> None:
         self.get_logger().info('Started Inference')
@@ -34,10 +36,9 @@ def start(self) -> None:
         self.get_logger().info('Finished Inference')
 
     def _prepare_dataset(self) -> Tuple[DataLoader, Tuple]:
-        raw_dataset = self.repository.get_all_data()
-        equalized_dataset = self.data_preprocessor.equalize_dataset_dimensions(raw_dataset)
-        inference_dataset = self.data_preprocessor.get_dataloader(equalized_dataset)
-        data_dimensions = self.data_preprocessor.extract_data_dimensions(equalized_dataset)
+        dataset = GraphDataset(self.data_path, test_mode=self.test_mode)
+        inference_dataset = self.data_preprocessor.get_dataloader(dataset)
+        data_dimensions = self.data_preprocessor.extract_data_dimensions(dataset)
         return inference_dataset, data_dimensions
 
     @staticmethod
diff --git a/message_passing_nn/utils/array_operations.cpp b/message_passing_nn/utils/array_operations.cpp
deleted file mode 100644
index 1a52e18..0000000
--- a/message_passing_nn/utils/array_operations.cpp
+++ /dev/null
@@ -1,38 +0,0 @@
-#include <torch/extension.h>
-#include "../utils/array_operations.h"
-
-std::vector<int> find_nonzero_elements(const torch::Tensor& tensor){
-  std::vector<int> vector;
-  for(int index=0; index<tensor.sizes()[0]; index++){
-    if(tensor[index].item<int>()!=0) {
-       vector.push_back(index);
-      }
-  }
-  return vector;
-}
-
-int find_index_by_value(const std::vector<int>& vector, const int& value){
-  auto vector_size = static_cast<int>(vector.size());
-  for (int index = 0; index < vector_size; ++index) {
-    if (vector[index]==value) return index;
-  }
-  return -1;
-}
-
-std::vector<int> remove_element_by_index_from_vector(std::vector<int>& vector, int& index){
-  std::vector<int> final_vector;
-  if (index==0) {
-    final_vector.assign(vector.begin() + 1, vector.end());
-  } else if ((index<static_cast<int>(vector.size()))){
-    std::vector<int> first_vector;
-    std::vector<int> second_vector;
-    first_vector.assign(vector.begin(), vector.begin() + index);
-    second_vector.assign(vector.begin() + index + 1, vector.end());
-    final_vector.reserve( first_vector.size() + second_vector.size() ); 
-    final_vector.insert( final_vector.end(), first_vector.begin(), first_vector.end() );
-    final_vector.insert( final_vector.end(), second_vector.begin(), second_vector.end() );
-  } else {
-    final_vector.assign(vector.begin(), vector.end() - 1);
-  }
-  return final_vector; 
-}
\ No newline at end of file
diff --git a/message_passing_nn/utils/array_operations.h b/message_passing_nn/utils/array_operations.h
deleted file mode 100644
index b4cb144..0000000
--- a/message_passing_nn/utils/array_operations.h
+++ /dev/null
@@ -1,10 +0,0 @@
-#ifndef ARRAY_OPERATIONS_H
-#define ARRAY_OPERATIONS_H
-
-#include <torch/extension.h>
-
-std::vector<int> find_nonzero_elements(const torch::Tensor& tensor);
-int find_index_by_value(const std::vector<int>& vector, const int& value);
-std::vector<int> remove_element_by_index_from_vector(std::vector<int>& vector, int& index);
-
-#endif
\ No newline at end of file
diff --git a/message_passing_nn/utils/derivatives.cpp b/message_passing_nn/utils/derivatives.cpp
index fa1a4d5..5832983 100644
--- a/message_passing_nn/utils/derivatives.cpp
+++ b/message_passing_nn/utils/derivatives.cpp
@@ -1,13 +1,13 @@
 #include <torch/extension.h>
 #include "../utils/derivatives.h"
 
-torch::Tensor d_sigmoid(torch::Tensor z) {
-  auto s = torch::sigmoid(z);
+at::Tensor d_sigmoid(at::Tensor z) {
+  auto s = at::sigmoid(z);
   return (1 - s) * s;
 }
 
-torch::Tensor d_relu_2d(torch::Tensor z) {
-  auto output = torch::zeros_like(z);
+at::Tensor d_relu_2d(at::Tensor z) {
+  auto output = at::zeros_like(z);
   for (int i = 0; i<z.sizes()[0]; i++) {
     for (int j = 0; j<z.sizes()[1]; j++) {
       if (z[i][j].item<float>() > 0.0) {
@@ -18,8 +18,8 @@ torch::Tensor d_relu_2d(torch::Tensor z) {
   return output;
 }
 
-torch::Tensor d_relu_4d(torch::Tensor z) {
-  auto output = torch::zeros_like(z);
+at::Tensor d_relu_4d(at::Tensor z) {
+  auto output = at::zeros_like(z);
   for (int i = 0; i<z.sizes()[0]; i++) {
     for (int j = 0; j<z.sizes()[1]; j++) {
       for (int k = 0; j<z.sizes()[1]; j++) {
diff --git a/message_passing_nn/utils/derivatives.h b/message_passing_nn/utils/derivatives.h
index 832b43b..c1d761b 100644
--- a/message_passing_nn/utils/derivatives.h
+++ b/message_passing_nn/utils/derivatives.h
@@ -3,8 +3,8 @@
 
 #include <torch/extension.h>
 
-torch::Tensor d_sigmoid(torch::Tensor z);
-torch::Tensor d_relu_2d(torch::Tensor z);
-torch::Tensor d_relu_4d(torch::Tensor z);
+at::Tensor d_sigmoid(at::Tensor z);
+at::Tensor d_relu_2d(at::Tensor z);
+at::Tensor d_relu_4d(at::Tensor z);
 
 #endif
\ No newline at end of file
diff --git a/message_passing_nn/utils/grid_search_parameters_parser.py b/message_passing_nn/utils/grid_search_parameters_parser.py
index 3e5f642..4894414 100644
--- a/message_passing_nn/utils/grid_search_parameters_parser.py
+++ b/message_passing_nn/utils/grid_search_parameters_parser.py
@@ -16,8 +16,6 @@ def get_grid_search_dictionary(self,
                                    loss_function_selection: str,
                                    optimizer_selection: str,
                                    batch_size: str,
-                                   maximum_number_of_features: str,
-                                   maximum_number_of_nodes: str,
                                    validation_split: str,
                                    test_split: str,
                                    time_steps: str,
@@ -28,8 +26,6 @@ def get_grid_search_dictionary(self,
             'loss_function': self._parse_string_selections(loss_function_selection),
             'optimizer': self._parse_string_selections(optimizer_selection),
             'batch_size': self._parse_integer_range(batch_size),
-            'maximum_number_of_features': self._parse_integer_range(maximum_number_of_features),
-            'maximum_number_of_nodes': self._parse_integer_range(maximum_number_of_nodes),
             'validation_split': self._parse_float_range(validation_split),
             'test_split': self._parse_float_range(test_split),
             'time_steps': self._parse_integer_range(time_steps),
diff --git a/message_passing_nn/utils/messages.cpp b/message_passing_nn/utils/messages.cpp
index db65c4e..39325d0 100644
--- a/message_passing_nn/utils/messages.cpp
+++ b/message_passing_nn/utils/messages.cpp
@@ -1,85 +1,56 @@
 #include <torch/extension.h>
 #include "../utils/messages.h"
-#include "../utils/array_operations.h"
 
-torch::Tensor compute_messages_from_neighbors(std::vector<int> all_neighbors,
-                                              const int& node_id,
-                                              const int& end_node_id,
-                                              const torch::Tensor& w_graph_neighbor_messages,
-                                              const torch::Tensor& messages_previous_step){
-  torch::Tensor messages_from_the_other_neighbors = torch::zeros_like({messages_previous_step[0][0]});
-  if (static_cast<int>(all_neighbors.size()) > 1) {
-          auto end_node_index = find_index_by_value(all_neighbors, end_node_id);
-          auto other_neighbors = remove_element_by_index_from_vector(all_neighbors, end_node_index);      
-          for (int neighbor: other_neighbors) {
-              messages_from_the_other_neighbors += torch::matmul(w_graph_neighbor_messages, messages_previous_step[neighbor][node_id]);
-          }
-        }
-  return messages_from_the_other_neighbors;
-}
-
-torch::Tensor get_messages_to_all_end_nodes(const int& node_id,
-                                            const torch::Tensor& w_graph_neighbor_messages,
-                                            const torch::Tensor& w_graph_node_features,
-                                            const torch::Tensor& adjacency_vector_of_specific_node,
-                                            const torch::Tensor& features_of_specific_node,
-                                            const torch::Tensor& messages_previous_step) {
-  auto all_neighbors = find_nonzero_elements(adjacency_vector_of_specific_node);
-  torch::Tensor new_messages_of_node = torch::zeros({messages_previous_step.sizes()[1], messages_previous_step.sizes()[2]});
-  for (int end_node_id: all_neighbors){
-      auto messages_from_the_other_neighbors = compute_messages_from_neighbors(all_neighbors,
-                                                                                node_id,
-                                                                                end_node_id,
-                                                                                w_graph_neighbor_messages,
-                                                                                messages_previous_step);
-      new_messages_of_node[end_node_id] = torch::add(torch::matmul(w_graph_node_features, features_of_specific_node), messages_from_the_other_neighbors);
-    }
-  return new_messages_of_node;
-}
-
-std::vector<torch::Tensor> compose_messages(
+std::vector<at::Tensor> compose_messages(
     const int& time_steps,
     const int& number_of_nodes,
     const int& number_of_node_features,
-    const torch::Tensor& w_graph_node_features,
-    const torch::Tensor& w_graph_neighbor_messages,
-    const torch::Tensor& node_features,
-    const torch::Tensor& adjacency_matrix,
-    const torch::Tensor& messages_init) {
+    const at::Tensor& w_graph_node_features,
+    const at::Tensor& w_graph_neighbor_messages,
+    const at::Tensor& base_messages,
+    const at::Tensor& all_neighbors,
+    const at::Tensor& messages_init) {
 
-  torch::Tensor new_messages = torch::zeros_like({messages_init});
-  torch::Tensor messages_previous_step;
+  auto new_messages = at::zeros_like({messages_init});
+  auto previous_messages = at::zeros_like({messages_init});
+  auto new_messages_of_node = at::zeros({previous_messages.sizes()[1], previous_messages.sizes()[2]});
   for (int time_step = 0; time_step<time_steps; time_step++) {
-    messages_previous_step = new_messages;
-    new_messages = torch::zeros_like({messages_init});
-
-    for (int node_id = 0; node_id<number_of_nodes; node_id++) {
-      new_messages[node_id] = get_messages_to_all_end_nodes(node_id,
-                                                            w_graph_neighbor_messages,
-                                                            w_graph_node_features,
-                                                            adjacency_matrix[node_id],
-                                                            node_features[node_id],
-                                                            messages_previous_step);
+    auto base_neighbor_messages = at::matmul(w_graph_neighbor_messages, at::relu(new_messages));
+    std::swap(previous_messages, new_messages);
+    for (int node_id = 0; node_id < all_neighbors.size(0); node_id++) {
+      for (int end_node_index = 0; end_node_index < all_neighbors.size(1); end_node_index++){
+        auto end_node_id = all_neighbors[node_id][end_node_index].item<int>();
+        if (end_node_id >= 0) {
+          new_messages[node_id][end_node_id] += base_messages[node_id];
+          for (int neighbor_index = 0; neighbor_index < all_neighbors.size(1); neighbor_index++) {
+            auto neighbor = all_neighbors[node_id][neighbor_index].item<int>();
+            if (neighbor >= 0 && neighbor_index!=end_node_index) {
+              new_messages[node_id][end_node_id] += base_neighbor_messages[neighbor][node_id];
+            }
+          }
+        }
+      }
     }
   }
-  return {new_messages, messages_previous_step};
+  return {new_messages, previous_messages};
 }
 
-torch::Tensor encode_messages(
+at::Tensor encode_messages(
     const int& number_of_nodes,
-    const torch::Tensor& node_encoding_messages,
-    const torch::Tensor& u_graph_node_features,
-    const torch::Tensor& u_graph_neighbor_messages,
-    const torch::Tensor& node_features,
-    const torch::Tensor& adjacency_matrix,
-    const torch::Tensor& messages) {
+    const at::Tensor& node_encoding_messages,
+    const at::Tensor& u_graph_node_features,
+    const at::Tensor& u_graph_neighbor_messages,
+    const at::Tensor& node_features,
+    const at::Tensor& all_neighbors,
+    const at::Tensor& messages) {
 
-    std::vector<int> all_neighbors;
     for (int node_id = 0; node_id<number_of_nodes; node_id++) {
-      all_neighbors = find_nonzero_elements(adjacency_matrix[node_id]);
-      for (int end_node_id: all_neighbors){
-        node_encoding_messages[node_id] += torch::matmul(u_graph_neighbor_messages, messages[end_node_id][node_id]);
+      for (int end_node_index = 0; end_node_index<all_neighbors.sizes()[1]; end_node_index++){
+        auto end_node_id = all_neighbors[node_id][end_node_index].item<int64_t>();
+        if (end_node_id >= 0) {
+          node_encoding_messages[node_id] += at::matmul(u_graph_neighbor_messages, at::relu(messages[end_node_id][node_id]));
+        }
       }
     }
-    return torch::relu(torch::add(torch::matmul(u_graph_node_features, node_features), node_encoding_messages));
+    return at::relu(at::add(at::matmul(u_graph_node_features, node_features), node_encoding_messages));
   }
diff --git a/message_passing_nn/utils/messages.h b/message_passing_nn/utils/messages.h
index 34cfc1a..fdbd50d 100644
--- a/message_passing_nn/utils/messages.h
+++ b/message_passing_nn/utils/messages.h
@@ -3,36 +3,30 @@
 
 #include <torch/extension.h>
 
-torch::Tensor compute_messages_from_neighbors(std::vector<int> all_neighbors,
-                                              const int& node_id,
-                                              const int& end_node_id,
-                                              const torch::Tensor& w_graph_neighbor_messages,
-                                              const torch::Tensor& messages_previous_step);
+at::Tensor get_messages_to_all_end_nodes(const int& node_id,
+                                            const at::Tensor& w_graph_neighbor_messages,
+                                            const at::Tensor& w_graph_node_features,
+                                            const at::Tensor& all_neighbors,
+                                            const at::Tensor& features_of_specific_node,
+                                            at::Tensor& messages_previous_step);
 
-torch::Tensor get_messages_to_all_end_nodes(const int& node_id,
-                                            const torch::Tensor& w_graph_neighbor_messages,
-                                            const torch::Tensor& w_graph_node_features,
-                                            const torch::Tensor& adjacency_vector_of_specific_node,
-                                            const torch::Tensor& features_of_specific_node,
-                                            const torch::Tensor& messages_previous_step);
-
-std::vector<torch::Tensor> compose_messages(
+std::vector<at::Tensor> compose_messages(
     const int& time_steps,
-    const int& number_of_nodes,
+    const int& number_of_true_nodes,
     const int& number_of_node_features,
-    const torch::Tensor& w_graph_node_features,
-    const torch::Tensor& w_graph_neighbor_messages,
-    const torch::Tensor& node_features,
-    const torch::Tensor& adjacency_matrix,
-    const torch::Tensor& messages_init);
+    const at::Tensor& w_graph_node_features,
+    const at::Tensor& w_graph_neighbor_messages,
+    const at::Tensor& base_messages,
+    const at::Tensor& all_neighbors,
+    const at::Tensor& messages_init);
 
-torch::Tensor encode_messages(
-    const int& number_of_nodes,
-    const torch::Tensor& node_encoding_messages,
-    const torch::Tensor& u_graph_node_features,
-    const torch::Tensor& u_graph_neighbor_messages,
-    const torch::Tensor& node_features,
-    const torch::Tensor& adjacency_matrix,
-    const torch::Tensor& messages);
+at::Tensor encode_messages(
+    const int& number_of_true_nodes,
+    const at::Tensor& node_encoding_messages,
+    const at::Tensor& u_graph_node_features,
+    const at::Tensor& u_graph_neighbor_messages,
+    const at::Tensor& node_features,
+    const at::Tensor& all_neighbors,
+    const at::Tensor& messages);
 
 #endif
\ No newline at end of file
diff --git a/message_passing_nn/utils/models.py b/message_passing_nn/utils/models.py
index b50bcd0..434350b 100644
--- a/message_passing_nn/utils/models.py
+++ b/message_passing_nn/utils/models.py
@@ -1,6 +1,5 @@
-from message_passing_nn.graph import rnn_encoder, gru_encoder
+from message_passing_nn.graph import rnn_encoder
 
 models = {
     "RNN": rnn_encoder.RNNEncoder,
-    "GRU": gru_encoder.GRUEncoder,
 }
diff --git a/parameters/grid-search-parameters.sh b/parameters/grid-search-parameters.sh
index 803fd00..f8b1b8c 100644
--- a/parameters/grid-search-parameters.sh
+++ b/parameters/grid-search-parameters.sh
@@ -1,18 +1,15 @@
 #!/usr/bin/env bash
-export DATASET_NAME='direct_neighbour_full_labels'
+export DATASET_NAME='sample-dataset'
 export DATA_DIRECTORY='data/'
 export MODEL_DIRECTORY='model_checkpoints'
 export RESULTS_DIRECTORY='results_grid_search'
 export MODEL='RNN'
 export DEVICE='cpu'
-export CPU_MULTIPROCESSING='True'
-export EPOCHS='400'
+export EPOCHS='10'
 export LOSS_FUNCTION='MSE'
 export OPTIMIZER='Adagrad'
 export BATCH_SIZE='100'
-export MAXIMUM_NUMBER_OF_FEATURES='-1'
-export MAXIMUM_NUMBER_OF_NODES='-1'
 export VALIDATION_SPLIT='0.2'
 export TEST_SPLIT='0.1'
 export TIME_STEPS='1'
-export VALIDATION_PERIOD='20'
+export VALIDATION_PERIOD='5'
diff --git a/parameters/inference-parameters.sh b/parameters/inference-parameters.sh
index 98e7501..d68f268 100644
--- a/parameters/inference-parameters.sh
+++ b/parameters/inference-parameters.sh
@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
-export DATASET_NAME='direct_neighbour_full_labels_test'
-export DATA_DIRECTORY='data/'
-export MODEL_DIRECTORY='model_checkpoints/configuration&id__model&RNN__epochs&400__loss_function&MSE__optimizer&Adagrad__batch_size&40__maximum_number_of_features&-1__maximum_number_of_nodes&-1__validation_split&0.2__test_split&0.1__time_steps&1__validation_period&40/Epoch_400_model_state_dictionary.pth'
+export DATASET_NAME='direct_neighbour'
+export DATA_DIRECTORY='data/the_100/'
+export MODEL_DIRECTORY='model_checkpoints/configuration&id__model&RNN__epochs&1000__loss_function&MSE__optimizer&Adagrad__batch_size&100__validation_split&0.2__test_split&0.1__time_steps&1__validation_period&5/Epoch_625_model_state_dictionary.pth'
 export RESULTS_DIRECTORY='results_inference'
 export MODEL='RNN'
 export DEVICE='cpu'
diff --git a/remove-containers-and-images.sh b/remove-containers-and-images.sh
deleted file mode 100644
index 2d5dfdb..0000000
--- a/remove-containers-and-images.sh
+++ /dev/null
@@ -1,8 +0,0 @@
-#!/bin/bash
-
-echo "Removing previous containers and images for message-passing-nn"
-docker rm -f message-passing-nn-container
-docker rmi message-passing-nn
-#echo "Removing stopped containers and untagged images"
-#docker container rm $(docker container ls -aq)
-#docker rmi $(docker images | grep "^<none>" | awk "{print $3}")
diff --git a/requirements.txt b/requirements.txt
index 286dc3e..eda8678 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 click
 numpy==1.17.4
 pandas==1.0.3
-tox==3.14.3
-torch==1.5.0
\ No newline at end of file
+torch==1.5.0
+tqdm
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 2d1d514..90165d1 100644
--- a/setup.py
+++ b/setup.py
@@ -1,12 +1,12 @@
 from setuptools import setup, find_packages, Extension
-from torch.utils import cpp_extension
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
 
 setup(
     name='message-passing-nn',
-    version='1.5.1',
+    version='1.6.0',
     packages=find_packages(exclude=["tests"]),
     url='https://github.com/kovanostra/message-passing-nn',
-    download_url='https://github.com/kovanostra/message-passing-nn/archive/1.5.0.tar.gz',
+    download_url='https://github.com/kovanostra/message-passing-nn/archive/1.6.0.tar.gz',
     keywords=['MESSAGE PASSING', 'NEURAL NETWORK', 'RNN', 'GRU'],
     license='MIT',
     author='Michail Kovanis',
@@ -15,7 +15,8 @@
         'click',
         'numpy==1.17.4',
         'pandas==1.0.3',
-        'torch==1.5.0'
+        'torch==1.5.0',
+        'tqdm'
     ],
     entry_points={
         'console_scripts': [
@@ -23,12 +24,15 @@
         ],
     },
     ext_modules=[
-        cpp_extension.CppExtension('rnn_encoder_cpp',
-                                   sources=['message_passing_nn/graph/rnn_encoder.cpp',
-                                            'message_passing_nn/utils/messages.cpp',
-                                            'message_passing_nn/utils/derivatives.cpp',
-                                            'message_passing_nn/utils/array_operations.cpp'])],
-    cmdclass={'build_ext': cpp_extension.BuildExtension},
+        CppExtension('rnn_encoder_cpp',
+                     sources=['message_passing_nn/graph/rnn_encoder.cpp',
+                              'message_passing_nn/utils/messages.cpp',
+                              'message_passing_nn/utils/derivatives.cpp']),
+        CUDAExtension('rnn_encoder_cuda_cpp',
+                      sources=['message_passing_nn/graph/rnn_encoder_cuda.cpp',
+                               'message_passing_nn/graph/rnn_encoder_cuda_kernel.cu',
+                               'message_passing_nn/utils/derivatives.cpp'])],
+    cmdclass={'build_ext': BuildExtension},
     classifiers=[
         'Development Status :: 5 - Production/Stable',
         'Intended Audience :: Science/Research',
diff --git a/setup_cpu.py b/setup_cpu.py
new file mode 100644
index 0000000..3d24166
--- /dev/null
+++ b/setup_cpu.py
@@ -0,0 +1,41 @@
+from setuptools import setup, find_packages, Extension
+from torch.utils.cpp_extension import BuildExtension, CppExtension
+
+setup(
+    name='message-passing-nn',
+    version='1.6.0',
+    packages=find_packages(exclude=["tests"]),
+    url='https://github.com/kovanostra/message-passing-nn',
+    download_url='https://github.com/kovanostra/message-passing-nn/archive/1.6.0.tar.gz',
+    keywords=['MESSAGE PASSING', 'NEURAL NETWORK', 'RNN', 'GRU'],
+    license='MIT',
+    author='Michail Kovanis',
+    description='A message passing neural network with RNN or GRU units',
+    install_requires=[
+        'click',
+        'numpy==1.17.4',
+        'pandas==1.0.3',
+        'torch==1.5.0',
+        'tqdm'
+    ],
+    entry_points={
+        'console_scripts': [
+            'message-passing-nn = message_passing_nn.cli:main'
+        ],
+    },
+    ext_modules=[
+        CppExtension('rnn_encoder_cpp',
+                     sources=['message_passing_nn/graph/rnn_encoder.cpp',
+                              'message_passing_nn/utils/messages.cpp',
+                              'message_passing_nn/utils/derivatives.cpp'])],
+    cmdclass={'build_ext': BuildExtension},
+    classifiers=[
+        'Development Status :: 5 - Production/Stable',
+        'Intended Audience :: Science/Research',
+        'Intended Audience :: Developers',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Software Development :: Libraries :: Python Modules',
+        'License :: OSI Approved :: MIT License',
+        'Programming Language :: Python :: 3.7',
+    ],
+)
diff --git a/tox.ini b/tox.ini
deleted file mode 100644
index 79d5edb..0000000
--- a/tox.ini
+++ /dev/null
@@ -1,6 +0,0 @@
-[tox]
-envlist = python3.7.6
-
-[testenv]
-deps=torch
-commands = python -m unittest discover
\ No newline at end of file


Test patch:
diff --git a/tests/data/test_data_preprocessor.py b/tests/data/test_data_preprocessor.py
index 18dae74..ea62596 100644
--- a/tests/data/test_data_preprocessor.py
+++ b/tests/data/test_data_preprocessor.py
@@ -1,6 +1,7 @@
 from unittest import TestCase
 
 import torch as to
+from message_passing_nn.infrastructure.graph_dataset import GraphDataset
 
 from message_passing_nn.data.data_preprocessor import DataPreprocessor
 from tests.fixtures.matrices_and_vectors import BASE_GRAPH, BASE_GRAPH_NODE_FEATURES
@@ -14,13 +15,18 @@ def test_train_validation_test_split(self):
         # Given
         dataset_length = 10
         features = BASE_GRAPH_NODE_FEATURES
-        adjacency_matrix = BASE_GRAPH
+        all_neighbors = to.tensor([[1, 2, -1, -1],
+                                   [0, 2, -1, -1],
+                                   [0, 1, 3, -1],
+                                   [2, -1, -1, -1]])
         labels = BASE_GRAPH.view(-1)
-        raw_dataset = [(features, adjacency_matrix, labels) for i in range(dataset_length)]
+        dataset = GraphDataset("")
+        dataset.enable_test_mode()
+        dataset.dataset = [(features, all_neighbors, labels, i) for i in range(dataset_length)]
         train_validation_test_split_expected = [7, 2, 1]
 
         # When
-        train_validation_test_split = self.data_preprocessor.train_validation_test_split(raw_dataset,
+        train_validation_test_split = self.data_preprocessor.train_validation_test_split(dataset,
                                                                                          batch_size=1,
                                                                                          validation_split=0.2,
                                                                                          test_split=0.1)
@@ -29,37 +35,22 @@ def test_train_validation_test_split(self):
         # Then
         self.assertEqual(train_validation_test_split_expected, train_validation_test_split)
 
-    def test_equalize_sizes(self):
-        # Given
-        node_features_1 = to.ones((2, 5))
-        node_features_2 = to.ones((3, 5))
-        adjacency_matrix_1 = to.ones(2, 2)
-        adjacency_matrix_2 = to.ones(3, 3)
-        labels_1 = to.ones(15)
-        labels_2 = to.ones(20)
-        raw_dataset = [(node_features_1, adjacency_matrix_1, labels_1),
-                       (node_features_2, adjacency_matrix_2, labels_2)]
-
-        # When
-        equalized_dataset = self.data_preprocessor.equalize_dataset_dimensions(raw_dataset, -1, -1)
-        node_features_1, adjacency_matrix_1, labels_1 = equalized_dataset[0]
-
-        # Then
-        self.assertEqual(node_features_1.size(), node_features_2.size())
-        self.assertEqual(adjacency_matrix_1.size(), adjacency_matrix_2.size())
-        self.assertEqual(labels_1.size(), labels_2.size())
-
     def test_extract_data_dimensions(self):
         # Given
         dataset_length = 1
         features = BASE_GRAPH_NODE_FEATURES
-        adjacency_matrix = BASE_GRAPH
+        all_neighbors = to.tensor([[1, 2, -1, -1],
+                                   [0, 2, -1, -1],
+                                   [0, 1, 3, -1],
+                                   [2, -1, -1, -1]])
         labels = BASE_GRAPH.view(-1)
-        raw_dataset = [(features, adjacency_matrix, labels) for i in range(dataset_length)]
-        data_dimensions_expected = (features.size(), adjacency_matrix.size(), labels.size())
+        dataset = GraphDataset("")
+        dataset.enable_test_mode()
+        dataset.dataset = [(features, all_neighbors, labels, i) for i in range(dataset_length)]
+        data_dimensions_expected = (features.size(), labels.size())
 
         # When
-        data_dimensions = self.data_preprocessor.extract_data_dimensions(raw_dataset)
+        data_dimensions = self.data_preprocessor.extract_data_dimensions(dataset)
 
         # Then
         self.assertEqual(data_dimensions_expected, data_dimensions)
diff --git a/tests/fixtures/matrices_and_vectors.py b/tests/fixtures/matrices_and_vectors.py
index 8ab6ef2..cca9027 100644
--- a/tests/fixtures/matrices_and_vectors.py
+++ b/tests/fixtures/matrices_and_vectors.py
@@ -1,11 +1,6 @@
 import numpy as np
 import torch as to
 
-NOT_SYMMETRIC_GRAPH = to.tensor([[0, 1, 1, 1],
-                                 [1, 0, 1, 0],
-                                 [1, 1, 0, 1],
-                                 [0, 0, 1, 0]]).float()
-
 BASE_GRAPH = to.tensor([[0, 1, 1, 0],
                         [1, 0, 1, 0],
                         [1, 1, 0, 1],
diff --git a/tests/graph/test_gru_encoder.py b/tests/graph/test_gru_encoder.py
deleted file mode 100644
index 1563a76..0000000
--- a/tests/graph/test_gru_encoder.py
+++ /dev/null
@@ -1,81 +0,0 @@
-from unittest import TestCase
-
-import torch as to
-from torch import nn
-
-from message_passing_nn.graph.gru_encoder import GRUEncoder
-from tests.fixtures.matrices_and_vectors import BASE_GRAPH, BASE_GRAPH_NODE_FEATURES, \
-    BASE_UNITY_MATRIX_TENSOR, MULTIPLICATION_FACTOR, BASE_B_VECTOR
-
-
-class TestGRUEncoder(TestCase):
-
-    def setUp(self) -> None:
-        number_of_nodes = BASE_GRAPH.size()[0]
-        number_of_node_features = BASE_GRAPH_NODE_FEATURES.size()[1]
-        self.graph_encoder = GRUEncoder.of(time_steps=1,
-                                           number_of_nodes=number_of_nodes,
-                                           number_of_node_features=number_of_node_features,
-                                           fully_connected_layer_input_size=number_of_nodes * number_of_node_features,
-                                           fully_connected_layer_output_size=number_of_nodes ** 2)
-        self.graph_encoder.w_gru_update_gate_features = nn.Parameter(MULTIPLICATION_FACTOR * BASE_UNITY_MATRIX_TENSOR,
-                                                                     requires_grad=False)
-        self.graph_encoder.w_gru_forget_gate_features = nn.Parameter(MULTIPLICATION_FACTOR * BASE_UNITY_MATRIX_TENSOR,
-                                                                     requires_grad=False)
-        self.graph_encoder.w_gru_current_memory_message_features = nn.Parameter(MULTIPLICATION_FACTOR * BASE_UNITY_MATRIX_TENSOR,
-                                                                                requires_grad=False)
-        self.graph_encoder.u_gru_update_gate = nn.Parameter(MULTIPLICATION_FACTOR * BASE_UNITY_MATRIX_TENSOR,
-                                                            requires_grad=False)
-        self.graph_encoder.u_gru_forget_gate = nn.Parameter(MULTIPLICATION_FACTOR * BASE_UNITY_MATRIX_TENSOR,
-                                                            requires_grad=False)
-        self.graph_encoder.u_gru_current_memory_message = nn.Parameter(MULTIPLICATION_FACTOR * BASE_UNITY_MATRIX_TENSOR,
-                                                                       requires_grad=False)
-        self.graph_encoder.b_gru_update_gate = nn.Parameter(MULTIPLICATION_FACTOR * BASE_B_VECTOR,
-                                                            requires_grad=False)
-        self.graph_encoder.b_gru_forget_gate = nn.Parameter(MULTIPLICATION_FACTOR * BASE_B_VECTOR,
-                                                            requires_grad=False)
-        self.graph_encoder.b_gru_current_memory_message = nn.Parameter(MULTIPLICATION_FACTOR * BASE_B_VECTOR,
-                                                                       requires_grad=False)
-        self.graph_encoder.u_graph_node_features = nn.Parameter(MULTIPLICATION_FACTOR * BASE_UNITY_MATRIX_TENSOR,
-                                                                requires_grad=False)
-        self.graph_encoder.u_graph_neighbor_messages = nn.Parameter(MULTIPLICATION_FACTOR * BASE_UNITY_MATRIX_TENSOR,
-                                                                    requires_grad=False)
-        self.graph_encoder.linear.weight = to.nn.Parameter(to.ones(number_of_nodes ** 2,
-                                                                   number_of_nodes * number_of_node_features),
-                                                           requires_grad=False).float()
-        self.graph_encoder.linear.bias = to.nn.Parameter(2 * to.ones(number_of_nodes ** 2), requires_grad=False).float()
-
-    def test_forward(self):
-        # Given
-        batch_size = 1
-        features = BASE_GRAPH_NODE_FEATURES.view(1, BASE_GRAPH.size()[0], BASE_GRAPH_NODE_FEATURES.size()[1])
-        adjacency_matrix = BASE_GRAPH.view(1, BASE_GRAPH.size()[0], BASE_GRAPH.size()[1],)
-        outputs_expected = to.tensor([[0.98733, 0.98733, 0.98733, 0.98733, 0.98733, 0.98733, 0.98733, 0.98733, 0.98733,
-                                       0.98733, 0.98733, 0.98733, 0.98733, 0.98733, 0.98733, 0.98733]])
-
-        # When
-        outputs = self.graph_encoder.forward(features, adjacency_matrix, batch_size)
-
-        # Then
-        self.assertTrue(to.allclose(outputs_expected, outputs))
-
-    def test_encode_graph_returns_the_expected_encoding_for_a_node_after_one_time_step(self):
-        # Give
-        node = 0
-        node_encoding_expected = to.tensor([[0.3909883, 0.3909883]])
-
-        # When
-        node_encoding = self.graph_encoder.encode(BASE_GRAPH_NODE_FEATURES, BASE_GRAPH)[node]
-
-        # Then
-        self.assertTrue(to.allclose(node_encoding_expected, node_encoding))
-
-    def test_encode_graph_returns_the_expected_shape(self):
-        # Given
-        encoded_graph_shape_expected = list(BASE_GRAPH_NODE_FEATURES.shape)
-
-        # When
-        encoded_graph_shape = self.graph_encoder.encode(BASE_GRAPH_NODE_FEATURES, BASE_GRAPH).shape
-
-        # Then
-        self.assertEqual(encoded_graph_shape_expected, list(encoded_graph_shape))
diff --git a/tests/graph/test_node.py b/tests/graph/test_node.py
deleted file mode 100644
index fe9cb89..0000000
--- a/tests/graph/test_node.py
+++ /dev/null
@@ -1,67 +0,0 @@
-from unittest import TestCase
-
-import numpy as np
-
-from message_passing_nn.graph.node import Node
-from tests.fixtures.matrices_and_vectors import BASE_GRAPH_NODE_FEATURES, BASE_GRAPH, NOT_SYMMETRIC_GRAPH
-
-
-class TestNode(TestCase):
-    def setUp(self) -> None:
-        all_node_features = BASE_GRAPH_NODE_FEATURES
-        labels = BASE_GRAPH
-        self.node = Node(all_node_features, labels, node_id=2)
-
-    def test_is_matrix_symmetric(self):
-        # Given
-        symmetric_matrix = BASE_GRAPH
-
-        # When
-        symmetry = Node.is_adjacency_matrix_symmetric(symmetric_matrix)
-
-        # Then
-        self.assertTrue(symmetry)
-
-    def test_is_matrix_not_symmetric(self):
-        # Given
-        not_symmetric_matrix = NOT_SYMMETRIC_GRAPH
-
-        # When
-        symmetry = Node.is_adjacency_matrix_symmetric(not_symmetric_matrix)
-
-        # Then
-        self.assertFalse(symmetry)
-
-    def test_node_neighbors_symmetric(self):
-        # Given
-        neighbors_expected = np.array([3])
-
-        # When
-        neighbors = self.node.neighbors
-
-        # Then
-        self.assertTrue(np.array_equal(neighbors_expected, neighbors))
-
-    def test_node_neighbors_not_symmetric(self):
-        # Given
-        all_node_features = BASE_GRAPH_NODE_FEATURES
-        labels = NOT_SYMMETRIC_GRAPH
-        self.node = Node(all_node_features, labels, node_id=2)
-        neighbors_expected = np.array([0, 1, 3])
-
-        # When
-        neighbors = self.node.neighbors
-
-        # Then
-        self.assertTrue(np.array_equal(neighbors_expected, neighbors))
-
-    def test_get_start_node_neighbors_without_end_node(self):
-        # Given
-        start_node_neighbors_expected = np.array([1, 3])
-        end_node_id = 0
-
-        # When
-        start_node_neighbors = self.node.get_start_node_neighbors_without_end_node(end_node_id)[0]
-
-        # Then
-        self.assertTrue(np.array_equal(start_node_neighbors_expected, start_node_neighbors))
diff --git a/tests/graph/test_rnn_encoder.py b/tests/graph/test_rnn_encoder.py
index ba03c1e..12ca0c5 100644
--- a/tests/graph/test_rnn_encoder.py
+++ b/tests/graph/test_rnn_encoder.py
@@ -24,10 +24,10 @@ def setUp(self) -> None:
                                         fully_connected_layer_input_size=self.fully_connected_layer_input_size,
                                         fully_connected_layer_output_size=self.fully_connected_layer_output_size)
         self.graph_encoder.w_graph_node_features = nn.Parameter(
-            MULTIPLICATION_FACTOR * (to.ones((self.number_of_node_features, self.number_of_node_features))),
+            MULTIPLICATION_FACTOR * (to.ones((self.number_of_nodes, self.number_of_nodes))),
             requires_grad=False)
         self.graph_encoder.w_graph_neighbor_messages = nn.Parameter(
-            MULTIPLICATION_FACTOR * to.ones((self.number_of_node_features, self.number_of_node_features)),
+            MULTIPLICATION_FACTOR * to.ones((self.number_of_nodes, self.number_of_nodes)),
             requires_grad=False)
         self.graph_encoder.u_graph_node_features = nn.Parameter(
             MULTIPLICATION_FACTOR * to.ones((self.number_of_nodes, self.number_of_nodes)),
@@ -48,16 +48,20 @@ def test_encode_graph_returns_the_expected_encoding_for_a_node_after_one_time_st
         node = 0
         time_steps = 1
         batch_size = 1
-        node_encoding_expected = to.tensor([0.5400, 0.4900])
+        all_neighbors = to.tensor([[[1, 2, -1, -1],
+                                    [0, 2, -1, -1],
+                                    [0, 1, 3, -1],
+                                    [2, -1, -1, -1]]])
+        node_encoding_expected = to.tensor([0.6200, 0.5700])
 
         # When
-        _, _, encodings, _, _ = rnn_cpp.forward(time_steps,
-                                                self.number_of_nodes,
-                                                self.number_of_node_features,
-                                                self.fully_connected_layer_output_size,
-                                                batch_size,
+        _, _, encodings, _, _ = rnn_cpp.forward(to.tensor(time_steps),
+                                                to.tensor(self.number_of_nodes),
+                                                to.tensor(self.number_of_node_features),
+                                                to.tensor(self.fully_connected_layer_output_size),
+                                                to.tensor(batch_size),
                                                 BASE_GRAPH_NODE_FEATURES.unsqueeze(0),
-                                                BASE_GRAPH.unsqueeze(0),
+                                                all_neighbors,
                                                 self.graph_encoder.w_graph_node_features,
                                                 self.graph_encoder.w_graph_neighbor_messages,
                                                 self.graph_encoder.u_graph_node_features,
@@ -73,16 +77,20 @@ def test_encode_graph_returns_the_expected_shape(self):
         # Given
         time_steps = 1
         batch_size = 1
+        all_neighbors = to.tensor([[[1, 2, -1, -1],
+                                    [0, 2, -1, -1],
+                                    [0, 1, 3, -1],
+                                    [2, -1, -1, -1]]])
         encoded_graph_shape_expected = list(BASE_GRAPH_NODE_FEATURES.view(1, -1).shape)
 
         # When
-        _, _, encodings, _, _ = rnn_cpp.forward(time_steps,
-                                                self.number_of_nodes,
-                                                self.number_of_node_features,
-                                                self.fully_connected_layer_output_size,
-                                                batch_size,
+        _, _, encodings, _, _ = rnn_cpp.forward(to.tensor(time_steps),
+                                                to.tensor(self.number_of_nodes),
+                                                to.tensor(self.number_of_node_features),
+                                                to.tensor(self.fully_connected_layer_output_size),
+                                                to.tensor(batch_size),
                                                 BASE_GRAPH_NODE_FEATURES.unsqueeze(0),
-                                                BASE_GRAPH.unsqueeze(0),
+                                                all_neighbors,
                                                 self.graph_encoder.w_graph_node_features,
                                                 self.graph_encoder.w_graph_neighbor_messages,
                                                 self.graph_encoder.u_graph_node_features,
@@ -100,18 +108,23 @@ def test_get_the_expected_messages_from_a_node_after_one_time_step(self):
                                      self.number_of_node_features),
                                     device=self.device)
         node_expected = 0
-        messages_from_node_expected = to.tensor([[0.0, 0.0],
-                                                 [0.3, 0.3],
-                                                 [0.3, 0.3],
-                                                 [0.0, 0.0]])
+        all_neighbors = to.tensor([[1, 2, -1, -1],
+                                   [0, 2, -1, -1],
+                                   [0, 1, 3, -1],
+                                   [2, -1, -1, -1]])
+        messages_from_node_expected = to.tensor([[0.00, 0.00],
+                                                 [0.45, 0.40],
+                                                 [0.45, 0.40],
+                                                 [0.00, 0.00]])
+        base_messages = to.matmul(self.graph_encoder.w_graph_node_features, BASE_GRAPH_NODE_FEATURES)
         # When
         messages_from_node, _ = rnn_cpp.compose_messages(time_steps,
                                                          self.graph_encoder.number_of_nodes,
                                                          self.graph_encoder.number_of_node_features,
                                                          self.graph_encoder.w_graph_node_features,
                                                          self.graph_encoder.w_graph_neighbor_messages,
-                                                         BASE_GRAPH_NODE_FEATURES,
-                                                         BASE_GRAPH,
+                                                         base_messages,
+                                                         all_neighbors,
                                                          messages_initial)
 
         # Then
@@ -126,10 +139,10 @@ def test_forward_for_batch_size_one_and_two_steps(self):
                                     [1.0, 1.0],
                                     [2.0, 0.5],
                                     [0.5, 0.5]]])
-        adjacency_matrix = to.tensor([[[0.0, 1.0, 1.0, 0.0],
-                                       [1.0, 0.0, 1.0, 0.0],
-                                       [1.0, 1.0, 0.0, 1.0],
-                                       [0.0, 0.0, 1.0, 0.0]]])
+        all_neighbors_input = to.tensor([[[1, 2, -1, -1],
+                                          [0, 2, -1, -1],
+                                          [0, 1, 3, -1],
+                                          [2, -1, -1, -1]]])
 
         # Calculations
         # -> Pre-loop
@@ -149,7 +162,8 @@ def test_forward_for_batch_size_one_and_two_steps(self):
         # -> Step 0
         # Calculations
         index_pairs_expected = [[0, 2], [0, 1], [1, 2], [1, 0], [2, 1], [2, 3], [2, 0], [2, 3], [2, 0], [2, 1]]
-        index_pairs, messages_step_0_part_1, messages_step_0_part_2 = self._message_calculations(all_neighbors, batch,
+        index_pairs, messages_step_0_part_1, messages_step_0_part_2 = self._message_calculations(all_neighbors,
+                                                                                                 batch,
                                                                                                  messages_init,
                                                                                                  messages_step_0_part_1,
                                                                                                  messages_step_0_part_2,
@@ -163,13 +177,14 @@ def test_forward_for_batch_size_one_and_two_steps(self):
         # -> Step 1
         # Messages
         messages_step_0 = to.relu(to.add(messages_step_0_part_1, messages_step_0_part_2))
+        base_messages = self.graph_encoder.w_graph_node_features.matmul(node_features[0])
         messages_from_model_step_0, _ = rnn_cpp.compose_messages(1,
                                                                  self.graph_encoder.number_of_nodes,
                                                                  self.graph_encoder.number_of_node_features,
                                                                  self.graph_encoder.w_graph_node_features,
                                                                  self.graph_encoder.w_graph_neighbor_messages,
-                                                                 node_features[0],
-                                                                 adjacency_matrix[0],
+                                                                 base_messages,
+                                                                 all_neighbors_input[0],
                                                                  messages_init)
         self.assertTrue(messages_step_0.size() == messages_init.size())
         self.assertTrue(np.allclose(messages_step_0.numpy(), messages_from_model_step_0.numpy(), atol=1e-02))
@@ -196,24 +211,25 @@ def test_forward_for_batch_size_one_and_two_steps(self):
         self._assert_step_1_parameters_are_correct(index_pairs, index_pairs_expected, messages_step_1_part_1,
                                                    messages_step_1_part_2, number_of_nodes)
 
-        # -> Step 1
+        # -> Step 2
         # Messages
         messages_step_1 = to.relu(to.add(messages_step_1_part_1, messages_step_1_part_2))
+        base_messages = self.graph_encoder.w_graph_node_features.matmul(node_features[0])
         messages_from_model_step_1, a = rnn_cpp.compose_messages(2,
                                                                  self.graph_encoder.number_of_nodes,
                                                                  self.graph_encoder.number_of_node_features,
                                                                  self.graph_encoder.w_graph_node_features,
                                                                  self.graph_encoder.w_graph_neighbor_messages,
-                                                                 node_features[0],
-                                                                 adjacency_matrix[0],
+                                                                 base_messages,
+                                                                 all_neighbors_input[0],
                                                                  messages_init)
         self.assertTrue(messages_step_1.size() == messages_init.size())
         self.assertTrue(np.allclose(messages_step_1.numpy(), messages_from_model_step_1.numpy(), atol=1e-02))
         print("Passed second step assertions!")
 
         # -> Sum messages
-        index_pairs, index_pairs_expected, messages_summed = self._sum_messages(all_neighbors, batch, index_pairs,
-                                                                                index_pairs_expected, messages_step_1,
+        index_pairs, index_pairs_expected, messages_summed = self._sum_messages(all_neighbors, batch,
+                                                                                messages_step_1,
                                                                                 number_of_node_features,
                                                                                 number_of_nodes)
         self.assertTrue(messages_summed.size() == to.empty(number_of_nodes, number_of_node_features).size())
@@ -235,12 +251,12 @@ def test_forward_for_batch_size_one_and_two_steps(self):
         print("Passed outputs assertions!")
 
         # When
-        outputs = self.graph_encoder.forward(node_features, adjacency_matrix, batch_size)
+        outputs = self.graph_encoder.forward(node_features, all_neighbors_input, batch_size)
 
         # Then
         self.assertTrue(np.allclose(outputs_expected.numpy(), outputs.numpy(), atol=1e-02))
 
-    def _sum_messages(self, all_neighbors, batch, index_pairs, index_pairs_expected, messages_step_1,
+    def _sum_messages(self, all_neighbors, batch, messages_step_1,
                       number_of_node_features, number_of_nodes):
         messages_summed = to.zeros(number_of_nodes, number_of_node_features, device=self.device)
         index_pairs_expected = np.array([[0, 1], [0, 2], [1, 0], [1, 2], [2, 0], [2, 1], [2, 3], [3, 2]])
@@ -258,10 +274,10 @@ def _sum_messages(self, all_neighbors, batch, index_pairs, index_pairs_expected,
 
     def _assert_step_1_parameters_are_correct(self, index_pairs, index_pairs_expected, messages_step_1_part_1,
                                               messages_step_1_part_2, number_of_nodes):
-        self.assertTrue(to.allclose(messages_step_1_part_1[0, 1], to.tensor([0.30, 0.30])))
-        self.assertTrue(to.allclose(messages_step_1_part_1[1, 2], to.tensor([0.20, 0.20])))
-        self.assertTrue(to.allclose(messages_step_1_part_1[2, 3], to.tensor([0.25, 0.25])))
-        self.assertTrue(to.allclose(messages_step_1_part_1[3, 2], to.tensor([0.10, 0.10])))
+        self.assertTrue(to.allclose(messages_step_1_part_1[0, 1], to.tensor([0.45, 0.40])))
+        self.assertTrue(to.allclose(messages_step_1_part_1[1, 2], to.tensor([0.45, 0.40])))
+        self.assertTrue(to.allclose(messages_step_1_part_1[2, 3], to.tensor([0.45, 0.40])))
+        self.assertTrue(to.allclose(messages_step_1_part_1[3, 2], to.tensor([0.45, 0.40])))
         for node_id in range(number_of_nodes):
             for end_node_id in range(number_of_nodes):
                 if [node_id, end_node_id] in index_pairs_expected:
@@ -273,10 +289,10 @@ def _assert_step_1_parameters_are_correct(self, index_pairs, index_pairs_expecte
 
     def _assert_step_0_parameters_are_correct(self, index_pairs, index_pairs_expected, messages_step_0_part_1,
                                               messages_step_0_part_2):
-        self.assertTrue(to.allclose(messages_step_0_part_1[0, 1], to.tensor([0.30, 0.30])))
-        self.assertTrue(to.allclose(messages_step_0_part_1[1, 2], to.tensor([0.20, 0.20])))
-        self.assertTrue(to.allclose(messages_step_0_part_1[2, 3], to.tensor([0.25, 0.25])))
-        self.assertTrue(to.allclose(messages_step_0_part_1[3, 2], to.tensor([0.10, 0.10])))
+        self.assertTrue(to.allclose(messages_step_0_part_1[0, 1], to.tensor([0.45, 0.40])))
+        self.assertTrue(to.allclose(messages_step_0_part_1[1, 2], to.tensor([0.45, 0.40])))
+        self.assertTrue(to.allclose(messages_step_0_part_1[2, 3], to.tensor([0.45, 0.40])))
+        self.assertTrue(to.allclose(messages_step_0_part_1[3, 2], to.tensor([0.45, 0.40])))
         self.assertTrue(np.array_equal(index_pairs_expected, np.array(index_pairs)))
         self.assertTrue(to.allclose(messages_step_0_part_2, to.zeros(
             (self.number_of_nodes, self.number_of_nodes, self.number_of_node_features))))
@@ -284,17 +300,18 @@ def _assert_step_0_parameters_are_correct(self, index_pairs, index_pairs_expecte
     def _message_calculations(self, all_neighbors, batch, messages_init, messages_step_0_part_1, messages_step_0_part_2,
                               neighbors_slice, node_features, number_of_node_features, number_of_nodes):
         index_pairs = []
+        base_messages = self.graph_encoder.w_graph_node_features.matmul(node_features)
+        base_neighbor_messages = self.graph_encoder.w_graph_neighbor_messages.matmul(messages_init)
         for node_id in range(number_of_nodes):
+            for index in range(len(all_neighbors[batch + node_id])):
+                end_node_id = all_neighbors[batch + node_id][index].item()
+                messages_from_neighbors_step_0 = to.zeros(number_of_node_features, device=self.device)
+                messages_step_0_part_1[node_id][end_node_id] += base_messages[batch][node_id]
+                if batch + node_id < len(neighbors_slice):
+                    for node_to_sum in neighbors_slice[batch + node_id][index]:
+                        index_pairs.append([node_id, node_to_sum])
+                        messages_from_neighbors_step_0 += base_neighbor_messages[node_to_sum, node_id]
+                    messages_step_0_part_2[node_id, end_node_id] = messages_from_neighbors_step_0
             if batch + node_id < len(all_neighbors):
-                for index in range(len(all_neighbors[batch + node_id])):
-                    end_node_id = all_neighbors[batch + node_id][index].item()
-                    messages_from_neighbors_step_0 = to.zeros(number_of_node_features, device=self.device)
-                    messages_step_0_part_1[node_id, end_node_id] = self.graph_encoder.w_graph_node_features.matmul(
-                        node_features[batch][node_id])
-                    if batch + node_id < len(neighbors_slice):
-                        for node_to_sum in neighbors_slice[batch + node_id][index]:
-                            index_pairs.append([node_id, node_to_sum])
-                            messages_from_neighbors_step_0 += self.graph_encoder.w_graph_neighbor_messages.matmul(
-                                messages_init[node_to_sum, node_id])
-                        messages_step_0_part_2[node_id, end_node_id] = messages_from_neighbors_step_0
+                pass
         return index_pairs, messages_step_0_part_1, messages_step_0_part_2
diff --git a/tests/repository/__init__.py b/tests/infrastructure/__init__.py
similarity index 100%
rename from tests/repository/__init__.py
rename to tests/infrastructure/__init__.py
diff --git a/tests/repository/test_file_system_repository.py b/tests/infrastructure/test_file_system_repository.py
similarity index 58%
rename from tests/repository/test_file_system_repository.py
rename to tests/infrastructure/test_file_system_repository.py
index 67bac57..d37747b 100644
--- a/tests/repository/test_file_system_repository.py
+++ b/tests/infrastructure/test_file_system_repository.py
@@ -2,9 +2,7 @@
 from os import path
 from unittest import TestCase
 
-import torch as to
-
-from message_passing_nn.repository.file_system_repository import FileSystemRepository
+from message_passing_nn.infrastructure.file_system_repository import FileSystemRepository
 from tests.fixtures.matrices_and_vectors import BASE_GRAPH, BASE_GRAPH_NODE_FEATURES
 
 
@@ -37,22 +35,3 @@ def test_save(self):
         os.remove(filenames_expected[0])
         os.remove(filenames_expected[1])
         os.remove(filenames_expected[2])
-
-    def test_get_all_data(self):
-        # Given
-        features_expected = BASE_GRAPH_NODE_FEATURES
-        adjacency_matrix_expected = BASE_GRAPH
-        labels_expected = BASE_GRAPH.view(-1)
-        filenames_to_save = ['code_features.pickle', 'code_adjacency-matrix.pickle', 'code_labels.pickle']
-        self.file_system_repository.save(filenames_to_save[0], features_expected)
-        self.file_system_repository.save(filenames_to_save[1], adjacency_matrix_expected)
-        self.file_system_repository.save(filenames_to_save[2], labels_expected)
-        all_data_expected = [(features_expected, adjacency_matrix_expected, labels_expected)]
-
-        # When
-        all_data = self.file_system_repository.get_all_data()
-
-        # Then
-        self.assertTrue(to.allclose(all_data_expected[0][0], all_data[0][0]))
-        self.assertTrue(to.allclose(all_data_expected[0][1], all_data[0][1]))
-        self.assertTrue(to.allclose(all_data_expected[0][2], all_data[0][2]))
diff --git a/tests/model/test_inferencer.py b/tests/model/test_inferencer.py
index 67f1355..9a65ca7 100644
--- a/tests/model/test_inferencer.py
+++ b/tests/model/test_inferencer.py
@@ -1,6 +1,7 @@
 from unittest import TestCase
 
-import numpy as np
+import torch as to
+from message_passing_nn.infrastructure.graph_dataset import GraphDataset
 
 from message_passing_nn.data.data_preprocessor import DataPreprocessor
 from message_passing_nn.model import Inferencer
@@ -21,8 +22,15 @@ def test_do_inference(self):
                       number_of_node_features=data_dimensions[0][1],
                       fully_connected_layer_input_size=data_dimensions[1][0] * data_dimensions[0][1],
                       fully_connected_layer_output_size=data_dimensions[2][0])
-        raw_dataset = [(BASE_GRAPH_NODE_FEATURES, BASE_GRAPH, BASE_GRAPH.view(-1))]
-        inference_data, _, _ = DataPreprocessor().train_validation_test_split(raw_dataset, 1, 0.0, 0.0)
+        all_neighbors = to.tensor([[1, 2, -1, -1],
+                                   [0, 2, -1, -1],
+                                   [0, 1, 3, -1],
+                                   [2, -1, -1, -1]])
+        dataset = GraphDataset("")
+        dataset.enable_test_mode()
+        tag = 'tag'
+        dataset.dataset = [(BASE_GRAPH_NODE_FEATURES, all_neighbors, BASE_GRAPH.view(-1), tag)]
+        inference_data, _, _ = DataPreprocessor().train_validation_test_split(dataset, 1, 0.0, 0.0)
         output_label_pairs_expected = [BASE_GRAPH.view(-1), BASE_GRAPH.view(-1)]
 
         # When
diff --git a/tests/model/test_loader.py b/tests/model/test_loader.py
index b652ad5..7115f51 100644
--- a/tests/model/test_loader.py
+++ b/tests/model/test_loader.py
@@ -8,11 +8,11 @@ class TestLoader(TestCase):
     def test_load_model(self):
         # Given
         loader = Loader("RNN")
-        data_dimensions = ([4, 2], [4, 4], [16])
-        path_to_model = "tests/test_data/model-checkpoints-test/configuration&id__model&RNN__epochs&10__" \
-                        "batch_size&3__maximum_number_of_nodes&-1__maximum_number_of_features&-1__" \
-                        "validation_split&0.2__test_split&0.1__loss_function&MSE__optimizer&SGD__" \
-                        "time_steps&1__validation_period&5/Epoch_5_model_state_dictionary.pth"
+        data_dimensions = ([4, 2], [16])
+        path_to_model = "tests/test_data/model-checkpoints-test/configuration&id__model&" + \
+                        "RNN__epochs&10__loss_function&MSE__optimizer&Adagrad__batch_size&" + \
+                        "100__validation_split&0.2__test_split&0.1__time_steps&1__validation_period&" + \
+                        "5/Epoch_5_model_state_dictionary.pth"
 
         # When
         model = loader.load_model(data_dimensions, path_to_model)
diff --git a/tests/model/test_trainer.py b/tests/model/test_trainer.py
index 5a270d2..e3f17a1 100644
--- a/tests/model/test_trainer.py
+++ b/tests/model/test_trainer.py
@@ -1,11 +1,14 @@
 from unittest import TestCase
 
+import torch as to
+from message_passing_nn.infrastructure.graph_dataset import GraphDataset
+
 from message_passing_nn.data.data_preprocessor import DataPreprocessor
 from message_passing_nn.model.trainer import Trainer
 from tests.fixtures.matrices_and_vectors import BASE_GRAPH, BASE_GRAPH_NODE_FEATURES
 
 
-class TestModelTrainer(TestCase):
+class TestTrainer(TestCase):
     def setUp(self) -> None:
         time_steps = 1
         loss_function = "MSE"
@@ -23,7 +26,7 @@ def test_instantiate_attributes(self):
         # Given
         number_of_nodes = BASE_GRAPH.size()[0]
         number_of_node_features = BASE_GRAPH_NODE_FEATURES.size()[1]
-        data_dimensions = (BASE_GRAPH_NODE_FEATURES.size(), BASE_GRAPH.size(), BASE_GRAPH.view(-1).size())
+        data_dimensions = (BASE_GRAPH_NODE_FEATURES.size(), BASE_GRAPH.view(-1).size())
 
         # When
         self.model_trainer.instantiate_attributes(data_dimensions, self.configuration_dictionary)
@@ -36,11 +39,18 @@ def test_instantiate_attributes(self):
 
     def test_do_train(self):
         # Given
-        data_dimensions = (BASE_GRAPH_NODE_FEATURES.size(), BASE_GRAPH.size(), BASE_GRAPH.view(-1).size())
+        data_dimensions = (BASE_GRAPH_NODE_FEATURES.size(), BASE_GRAPH.view(-1).size())
         self.model_trainer.instantiate_attributes(data_dimensions,
                                                   self.configuration_dictionary)
-        raw_dataset = [(BASE_GRAPH_NODE_FEATURES, BASE_GRAPH, BASE_GRAPH.view(-1))]
-        training_data, _, _ = DataPreprocessor().train_validation_test_split(raw_dataset, 1, 0.0, 0.0)
+        all_neighbors = to.tensor([[1, 2, -1, -1],
+                                   [0, 2, -1, -1],
+                                   [0, 1, 3, -1],
+                                   [2, -1, -1, -1]])
+        dataset = GraphDataset("")
+        dataset.enable_test_mode()
+        tag = 'tag'
+        dataset.dataset = [(BASE_GRAPH_NODE_FEATURES, all_neighbors, BASE_GRAPH.view(-1), tag)]
+        training_data, _, _ = DataPreprocessor().train_validation_test_split(dataset, 1, 0.0, 0.0)
 
         # When
         training_loss = self.model_trainer.do_train(training_data=training_data, epoch=1)
@@ -50,11 +60,18 @@ def test_do_train(self):
 
     def test_do_evaluate(self):
         # Given
-        data_dimensions = (BASE_GRAPH_NODE_FEATURES.size(), BASE_GRAPH.size(), BASE_GRAPH.view(-1).size())
+        data_dimensions = (BASE_GRAPH_NODE_FEATURES.size(), BASE_GRAPH.view(-1).size())
         self.model_trainer.instantiate_attributes(data_dimensions,
                                                   self.configuration_dictionary)
-        raw_dataset = [(BASE_GRAPH_NODE_FEATURES, BASE_GRAPH, BASE_GRAPH.view(-1))]
-        training_data, _, _ = DataPreprocessor().train_validation_test_split(raw_dataset, 1, 0.0, 0.0)
+        all_neighbors = to.tensor([[1, 2, -1, -1],
+                                   [0, 2, -1, -1],
+                                   [0, 1, 3, -1],
+                                   [2, -1, -1, -1]])
+        dataset = GraphDataset("")
+        dataset.enable_test_mode()
+        tag = 'tag'
+        dataset.dataset = [(BASE_GRAPH_NODE_FEATURES, all_neighbors, BASE_GRAPH.view(-1), tag)]
+        training_data, _, _ = DataPreprocessor().train_validation_test_split(dataset, 1, 0.0, 0.0)
 
         # When
         validation_loss = self.model_trainer.do_evaluate(evaluation_data=training_data, epoch=1)
diff --git a/tests/test_data/model-checkpoints-test/configuration&id__model&RNN__epochs&10__batch_size&3__maximum_number_of_nodes&-1__maximum_number_of_features&-1__validation_split&0.2__test_split&0.1__loss_function&MSE__optimizer&SGD__time_steps&1__validation_period&5/Epoch_5_model_state_dictionary.pth b/tests/test_data/model-checkpoints-test/configuration&id__model&RNN__epochs&10__batch_size&3__maximum_number_of_nodes&-1__maximum_number_of_features&-1__validation_split&0.2__test_split&0.1__loss_function&MSE__optimizer&SGD__time_steps&1__validation_period&5/Epoch_5_model_state_dictionary.pth
deleted file mode 100644
index 0a19152..0000000
Binary files a/tests/test_data/model-checkpoints-test/configuration&id__model&RNN__epochs&10__batch_size&3__maximum_number_of_nodes&-1__maximum_number_of_features&-1__validation_split&0.2__test_split&0.1__loss_function&MSE__optimizer&SGD__time_steps&1__validation_period&5/Epoch_5_model_state_dictionary.pth and /dev/null differ
diff --git a/tests/test_data/model-checkpoints-test/configuration&id__model&RNN__epochs&10__loss_function&MSE__optimizer&Adagrad__batch_size&100__validation_split&0.2__test_split&0.1__time_steps&1__validation_period&5/Epoch_5_model_state_dictionary.pth b/tests/test_data/model-checkpoints-test/configuration&id__model&RNN__epochs&10__loss_function&MSE__optimizer&Adagrad__batch_size&100__validation_split&0.2__test_split&0.1__time_steps&1__validation_period&5/Epoch_5_model_state_dictionary.pth
new file mode 100644
index 0000000..3231dea
Binary files /dev/null and b/tests/test_data/model-checkpoints-test/configuration&id__model&RNN__epochs&10__loss_function&MSE__optimizer&Adagrad__batch_size&100__validation_split&0.2__test_split&0.1__time_steps&1__validation_period&5/Epoch_5_model_state_dictionary.pth differ
diff --git a/tests/usecase/test_grid_search.py b/tests/usecase/test_grid_search.py
index 63e1712..159d46a 100644
--- a/tests/usecase/test_grid_search.py
+++ b/tests/usecase/test_grid_search.py
@@ -4,7 +4,7 @@
 
 from message_passing_nn.data.data_preprocessor import DataPreprocessor
 from message_passing_nn.model.trainer import Trainer
-from message_passing_nn.repository.file_system_repository import FileSystemRepository
+from message_passing_nn.infrastructure.file_system_repository import FileSystemRepository
 from message_passing_nn.usecase.grid_search import GridSearch
 from message_passing_nn.utils.saver import Saver
 from tests.fixtures.matrices_and_vectors import BASE_GRAPH, BASE_GRAPH_NODE_FEATURES
@@ -20,21 +20,20 @@ def setUp(self) -> None:
         tests_model_directory = 'tests/model_checkpoints'
         tests_results_directory = 'tests/grid_search_results'
         device = "cpu"
-        cpu_multiprocessing = False
+        self.data_path = self.tests_data_directory + self.dataset + "/"
         self.repository = FileSystemRepository(self.tests_data_directory, self.dataset)
         self.data_preprocessor = DataPreprocessor()
-        self.model_trainer = Trainer(self.data_preprocessor, device, cpu_multiprocessing)
+        self.data_preprocessor.enable_test_mode()
+        self.model_trainer = Trainer(self.data_preprocessor, device)
         self.saver = Saver(tests_model_directory, tests_results_directory)
 
     def test_start_for_multiple_batches_of_the_same_size(self):
         # Given
         dataset_size = 6
         grid_search_dictionary = {
-            "model": ["GRU"],
+            "model": ["RNN"],
             "epochs": [10],
             "batch_size": [3],
-            "maximum_number_of_nodes": [-1],
-            "maximum_number_of_features": [-1],
             "validation_split": [0.2],
             "test_split": [0.1],
             "loss_function": ["MSE"],
@@ -42,11 +41,12 @@ def test_start_for_multiple_batches_of_the_same_size(self):
             "time_steps": [1],
             "validation_period": [5]
         }
-        grid_search = GridSearch(self.repository,
+        grid_search = GridSearch(self.data_path,
                                  self.data_preprocessor,
                                  self.model_trainer,
                                  grid_search_dictionary,
-                                 self.saver)
+                                 self.saver,
+                                 test_mode=True)
 
         adjacency_matrix_filenames, features_filenames, labels_filenames = self._save_test_data(dataset_size)
 
@@ -63,16 +63,6 @@ def test_start_for_multiple_batches_of_the_same_size(self):
         # Tear down
         self._remove_files(dataset_size, features_filenames, adjacency_matrix_filenames, labels_filenames)
 
-    def _save_test_data(self, dataset_size):
-        features_filenames = [str(i) + '_training_features' + '.pickle' for i in range(dataset_size)]
-        adjacency_matrix_filenames = [str(i) + '_training_adjacency-matrix' '.pickle' for i in range(dataset_size)]
-        labels_filenames = [str(i) + '_training_labels' '.pickle' for i in range(dataset_size)]
-        for i in range(dataset_size):
-            self.repository.save(features_filenames[i], self.features)
-            self.repository.save(adjacency_matrix_filenames[i], self.adjacency_matrix)
-            self.repository.save(labels_filenames[i], self.labels)
-        return adjacency_matrix_filenames, features_filenames, labels_filenames
-
     def test_start_for_multiple_batches_of_differing_size(self):
         # Given
         dataset_size = 5
@@ -80,8 +70,6 @@ def test_start_for_multiple_batches_of_differing_size(self):
             "model": ["RNN"],
             "epochs": [10],
             "batch_size": [3],
-            "maximum_number_of_nodes": [-1],
-            "maximum_number_of_features": [-1],
             "validation_split": [0.2],
             "test_split": [0.1],
             "loss_function": ["MSE"],
@@ -89,11 +77,12 @@ def test_start_for_multiple_batches_of_differing_size(self):
             "time_steps": [1],
             "validation_period": [5]
         }
-        grid_search = GridSearch(self.repository,
+        grid_search = GridSearch(self.data_path,
                                  self.data_preprocessor,
                                  self.model_trainer,
                                  grid_search_dictionary,
-                                 self.saver)
+                                 self.saver,
+                                 test_mode=True)
 
         adjacency_matrix_filenames, features_filenames, labels_filenames = self._save_test_data(dataset_size)
 
@@ -114,11 +103,9 @@ def test_start_a_grid_search(self):
         # Given
         dataset_size = 6
         grid_search_dictionary = {
-            "model": ["RNN", "GRU"],
+            "model": ["RNN"],
             "epochs": [10, 15],
             "batch_size": [3, 4],
-            "maximum_number_of_nodes": [-1],
-            "maximum_number_of_features": [-1],
             "validation_split": [0.2],
             "test_split": [0.1],
             "loss_function": ["MSE"],
@@ -126,11 +113,12 @@ def test_start_a_grid_search(self):
             "time_steps": [1],
             "validation_period": [5]
         }
-        grid_search = GridSearch(self.repository,
+        grid_search = GridSearch(self.data_path,
                                  self.data_preprocessor,
                                  self.model_trainer,
                                  grid_search_dictionary,
-                                 self.saver)
+                                 self.saver,
+                                 test_mode=True)
 
         adjacency_matrix_filenames, features_filenames, labels_filenames = self._save_test_data(dataset_size)
 
@@ -147,6 +135,16 @@ def test_start_a_grid_search(self):
         # Tear down
         self._remove_files(dataset_size, features_filenames, adjacency_matrix_filenames, labels_filenames)
 
+    def _save_test_data(self, dataset_size):
+        features_filenames = [str(i) + '_training_features' + '.pickle' for i in range(dataset_size)]
+        adjacency_matrix_filenames = [str(i) + '_training_adjacency-matrix' '.pickle' for i in range(dataset_size)]
+        labels_filenames = [str(i) + '_training_labels' '.pickle' for i in range(dataset_size)]
+        for i in range(dataset_size):
+            self.repository.save(features_filenames[i], self.features)
+            self.repository.save(adjacency_matrix_filenames[i], self.adjacency_matrix)
+            self.repository.save(labels_filenames[i], self.labels)
+        return adjacency_matrix_filenames, features_filenames, labels_filenames
+
     def _remove_files(self,
                       dataset_size: int,
                       features_filenames: List[str],
diff --git a/tests/usecase/test_inference.py b/tests/usecase/test_inference.py
index 44571a4..fc28d53 100644
--- a/tests/usecase/test_inference.py
+++ b/tests/usecase/test_inference.py
@@ -7,7 +7,7 @@
 
 from message_passing_nn.data.data_preprocessor import DataPreprocessor
 from message_passing_nn.model import Loader, Inferencer
-from message_passing_nn.repository.file_system_repository import FileSystemRepository
+from message_passing_nn.infrastructure.file_system_repository import FileSystemRepository
 from message_passing_nn.usecase import Inference
 from message_passing_nn.utils.saver import Saver
 
@@ -19,28 +19,33 @@ def test_start(self):
         features = to.ones(4, 2)
         adjacency_matrix = to.ones(4, 4)
         labels = to.ones(16)
-        dataset = 'training-test-data'
+        dataset = 'inference-test-data'
         tests_data_directory = 'tests/test_data/'
-        tests_model_directory = "tests/test_data/model-checkpoints-test/configuration&id__model&RNN__epochs&10__" \
-                                "batch_size&3__maximum_number_of_nodes&-1__maximum_number_of_features&-1__" \
-                                "validation_split&0.2__test_split&0.1__loss_function&MSE__optimizer&SGD__" \
-                                "time_steps&1__validation_period&5/Epoch_5_model_state_dictionary.pth"
+        tests_model_directory = "tests/test_data/model-checkpoints-test/configuration&id__model&" + \
+                                "RNN__epochs&10__loss_function&MSE__optimizer&Adagrad__batch_size&" + \
+                                "100__validation_split&0.2__test_split&0.1__time_steps&1__validation_period&" + \
+                                "5/Epoch_5_model_state_dictionary.pth"
         tests_results_directory = 'tests/results_inference'
         device = "cpu"
         repository = FileSystemRepository(tests_data_directory, dataset)
+        data_path = tests_data_directory + dataset + "/"
         data_preprocessor = DataPreprocessor()
+        data_preprocessor.enable_test_mode()
         loader = Loader("RNN")
         inferencer = Inferencer(data_preprocessor, device)
         saver = Saver(tests_model_directory, tests_results_directory)
-        inference = Inference(repository,
+        inference = Inference(data_path,
                               data_preprocessor,
                               loader,
                               inferencer,
-                              saver)
+                              saver,
+                              test_mode=True)
 
         adjacency_matrix_filenames, features_filenames, labels_filenames = self._save_test_data(adjacency_matrix,
-                                                                                                dataset_size, features,
-                                                                                                labels, repository)
+                                                                                                dataset_size,
+                                                                                                features,
+                                                                                                labels,
+                                                                                                repository)
 
         # When
         inference.start()
diff --git a/tests/utils/test_grid_search_parameters_parser.py b/tests/utils/test_grid_search_parameters_parser.py
index 8a0ec10..407b6d9 100644
--- a/tests/utils/test_grid_search_parameters_parser.py
+++ b/tests/utils/test_grid_search_parameters_parser.py
@@ -11,8 +11,6 @@ def test_get_grid_search_dictionary(self):
         loss_function_selection = "MSE&CrossEntropy"
         optimizer_selection = "SGD&Adam"
         batch_size = "5"
-        maximum_number_of_nodes = "-1"
-        maximum_number_of_features = "-1"
         validation_split = "0.2"
         test_split = "0.1&0.2&2"
         time_steps = "10"
@@ -23,8 +21,6 @@ def test_get_grid_search_dictionary(self):
             "loss_function": ["MSE", "CrossEntropy"],
             "optimizer": ["SGD", "Adam"],
             "batch_size": [5],
-            "maximum_number_of_nodes": [-1],
-            "maximum_number_of_features": [-1],
             "validation_split": [0.2],
             "test_split": [0.1, 0.2],
             "time_steps": [10],
@@ -38,8 +34,6 @@ def test_get_grid_search_dictionary(self):
             loss_function_selection,
             optimizer_selection,
             batch_size,
-            maximum_number_of_nodes,
-            maximum_number_of_features,
             validation_split,
             test_split,
             time_steps,