diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md
deleted file mode 100644
index c9b931a..0000000
--- a/.claude/CLAUDE.md
+++ /dev/null
@@ -1,131 +0,0 @@
-# CLAUDE.md - SimpleGNN Project Guide
-
-## Project Overview
-
-SimpleGNN is a PyTorch-based Graph Neural Network experimentation framework for benchmarking and developing GNN architectures. It supports classical message-passing GNNs (GCN, GIN, GAT, GATv2, GraphSAGE) and a proprietary ShareGNN variant with invariant-based layers. The framework handles graph classification, graph regression, and node classification tasks.
-
-## Repository Structure
-
-```
-SimpleGNN/
-├── src/                              # All source code
-│   ├── framework/                    # Training/evaluation orchestration
-│   │   ├── core.py                   # FrameworkMain - main entry point class
-│   │   ├── model_configuration.py    # Single config training: model init, train/eval loops
-│   │   ├── run_configuration.py      # Hyperparameter grid search combinations
-│   │   └── utils/
-│   │       ├── parameters.py         # Parameters class (all experiment hyperparams)
-│   │       ├── preprocessing.py      # Dataset loading, splits, label computation
-│   │       ├── evaluation.py         # Result analysis and visualization
-│   │       ├── configuration_checks.py  # YAML config validation
-│   │       └── data_sampling.py      # Batch sampling strategies
-│   ├── models/                       # GNN models and layers
-│   │   ├── model.py                  # GraphModel - main PyTorch model class
-│   │   ├── layers/
-│   │   │   ├── framework_layer.py    # Abstract base class for all custom layers
-│   │   │   ├── mpnn_classical/       # Classical GNN wrappers (gcn, gin, gat, gatv2, sage, pooling)
-│   │   │   ├── nn_standard/          # Standard layers (linear, activation, batchnorm, dropout, reshape)
-│   │   │   └── utils/                # LayerTypes enum, layer_loader
-│   │   └── ShareGNN/                 # Proprietary ShareGNN implementation
-│   │       ├── layers/               # inv_based_message_passing, inv_based_pooling, positional_encoding
-│   │       ├── preprocessing/        # ShareGNN-specific label/property preprocessing
-│   │       └── utils.py
-│   ├── datasets/                     # Dataset handling
-│   │   ├── graph_dataset.py          # GraphDataset (PyG InMemoryDataset)
-│   │   ├── graph_dataset_preprocessing.py  # Dataset-specific preprocessing
-│   │   ├── custom_datasets.py        # Dataset factory and registration
-│   │   ├── custom_benchmarks/        # Synthetic benchmarks (rings, snowflakes, strings)
-│   │   ├── splits/                   # Pre-computed train/val/test splits (JSON)
-│   │   └── utils/                    # Node/edge label utilities, graph functions
-│   └── utils/                        # General utilities (timer, path conversions)
-├── examples/                         # Example experiments with YAML configs
-│   ├── basic_example/                # GCN/GIN on MUTAG
-│   ├── basic_example_share_gnn/      # ShareGNN on MUTAG
-│   └── zinc/                         # ShareGNN on ZINC
-├── experiments/                      # Shell scripts for reproducing paper results
-│   └── base_paper/                   # TUDatasets, ZINC, QM9, ablation, synthetic benchmarks
-├── data/                             # Dataset storage (gitignored, auto-downloaded)
-├── results/                          # Experiment results (gitignored)
-└── docs/                             # Sphinx documentation
-```
-
-## Configuration System
-
-The framework uses a **three-tier YAML configuration**:
-
-1. **Main config** (`main.yml`): Datasets, task type, paths to model/hyperparameter configs, splits
-2. **Model config** (`models_*.yml`): Layer architecture as list of layer definitions (supports grid search via list of lists)
-3. **Hyperparameter config** (`parameters.yml`): Training params (optimizer, loss, lr, epochs, batch size, input features). Lists = grid search.
-
-Configs are validated by `src/framework/utils/configuration_checks.py` against mandatory parameter sets.
-
-## Running Experiments
-
-Entry point pattern (see `examples/*/main.py`):
-
-```python
-from pathlib import Path
-from framework.core import FrameworkMain
-
-experiment = FrameworkMain(Path('examples/basic_example/main.yml'))
-experiment.preprocessing(num_threads=1)        # Load data, generate labels/splits
-experiment.run_configurations(num_threads=-1)   # Grid search (-1 = all CPUs)
-experiment.evaluate_results()                   # Find best config on validation set
-experiment.run_best_configuration(num_threads=-1)  # Re-run best config
-experiment.evaluate_results(evaluate_best_model=True)  # Final test-set evaluation
-```
-
-Run from the `src/` directory (imports are relative to `src/`):
-```bash
-cd src && python -m examples.basic_example.main
-```
-
-Experiment shell scripts are in `experiments/base_paper/`.
-
-## Key Architecture Decisions
-
-- **GraphModel** (`src/models/model.py`): Sequential `nn.ModuleList` of layers built from YAML config. All layers extend `FrameworkLayer` base class.
-- **Layer types**: Defined in `LayerTypes` enum (`src/models/layers/utils/layer_types.py`). New layers must be registered there and in `layer_loader.py`.
-- **Tensor shapes**: Layers handle `(C, N, F)` for multi-channel/multi-head data or `(N, F)` for standard. C = channels/heads, N = nodes, F = features.
-- **LinearLayer modes**: `aggr_features` (standard), `aggr_channels` (aggregate across channels), `channel_wise` (independent per channel).
-- **ShareGNN layers** use node/edge labels and pairwise properties for invariant-based message passing (multi-head output).
-
-## Dependencies
-
-Core: PyTorch, PyTorch Geometric (torch_geometric), numpy, pandas, scikit-learn, networkx, joblib, pyyaml, matplotlib.
-Optional: ogb, rdkit (molecular datasets).
-
-Virtual environment is in `venv/`. Note: `requirements.txt` and `setup.py` have been removed from tracking.
-
-## Coding Conventions
-
-- **Classes**: PascalCase (`FrameworkMain`, `GraphDataset`, `GCNConv`)
-- **Functions/methods**: snake_case (`run_configurations`, `forward`, `load_preprocessed_data_and_parameters`)
-- **Constants**: UPPER_SNAKE_CASE (`MANDATORY_MAIN_CONFIG_PARAMS`)
-- **Private members**: underscore prefix (`_num_graph_nodes`)
-- **Imports**: Standard library first, then third-party (torch, numpy), then local modules. Local imports use dot-notation relative to `src/` (e.g., `from models.model import GraphModel`).
-- **Type hints**: Used in method signatures where present, not comprehensive across the codebase.
-- **Indentation**: 4 spaces.
-- **Docstrings**: Present on classes and key methods; not exhaustive. Don't add docstrings to code you didn't write.
-
-## Git Conventions
-
-- **Main branch**: `master`
-- **Commit messages**: Imperative mood, descriptive ("Refactor linear layer implementation to support 'aggr_channels' mode")
-- **No unit test suite**: Testing is done via integration experiments (example scripts and experiment shell scripts)
-- **Gitignored**: `data/`, `results/`, `*.csv`, `venv/`, `__pycache__/`, `.idea/`, `.auto-claude/`, `/.claude/`
-
-## Important Patterns
-
-- **Adding a new GNN layer**: Create wrapper class extending `FrameworkLayer` in `src/models/layers/mpnn_classical/`, add to `LayerTypes` enum, register in `layer_loader.py`, then reference in YAML model config.
-- **Adding a new dataset**: Implement preprocessing class in `graph_dataset_preprocessing.py`, register in `custom_datasets.py`, create split files in `datasets/splits/`.
-- **Grid search**: Use lists in YAML parameter files. The framework computes the cartesian product of all list-valued parameters.
-- **Parallel execution**: `joblib` handles parallel runs across splits/configs. `num_threads=-1` uses all CPUs.
-- **Results**: Saved as CSV per epoch per configuration in the results directory. Evaluation finds best config by validation metric.
-
-## Common Pitfalls
-
-- Always run from `src/` directory or ensure `src/` is on the Python path; imports are relative to it.
-- YAML model configs use `layer_type` string keys that must match `LayerTypes` enum values exactly.
-- ShareGNN preprocessing must run before training ShareGNN models (generates required node/edge properties).
-- The `precision` parameter (`float`/`double`) must be consistent between data and model; mismatches cause runtime errors.
diff --git a/src/datasets/graph_dataset.py b/src/datasets/graph_dataset.py
index 3f0f3b9..81cc724 100644
--- a/src/datasets/graph_dataset.py
+++ b/src/datasets/graph_dataset.py
@@ -11,10 +11,10 @@
 
 from datasets.utils.EdgeLabels import EdgeLabels
 from datasets.utils.NodeLabels import NodeLabels
-from src.datasets.graph_dataset_preprocessing import ZINCGraphDataPreprocessing, QMGraphDataPreprocessing, \
+from datasets.graph_dataset_preprocessing import ZINCGraphDataPreprocessing, QMGraphDataPreprocessing, \
     OGBGraphPropertyGraphDataPreprocessing, SubstructureBenchmarkPreprocessing, MergedGraphDataPreprocessing, \
     TUDatasetPreprocessing
-from src.utils.utils import load_graphs
+from utils.utils import load_graphs
 from torch_geometric.io import fs
 from torch_geometric.utils.convert import to_networkx
 from ogb.nodeproppred import PygNodePropPredDataset
diff --git a/src/framework/core.py b/src/framework/core.py
index 188a43e..93a075a 100644
--- a/src/framework/core.py
+++ b/src/framework/core.py
@@ -504,17 +504,12 @@ def collect_paths(main_configuration, model_configuration, dataset_configuration
 
 def copy_experiment_config(absolute_path, experiment_configuration, experiment_configuration_path,
                            graph_db_name):
+    import shutil
     results_path = experiment_configuration['paths']['results']
     if not results_path.joinpath(f"{graph_db_name}/config.yml"):
         source_path = Path(absolute_path).joinpath(experiment_configuration_path)
-        destination_path =results_path.joinpath(f"{graph_db_name}/config.yml")
-        # copy the config file to the results directory
-        # if linux
-        if os.name == 'posix':
-            os.system(f"cp {source_path} {destination_path}")
-        # if windows
-        elif os.name == 'nt':
-            os.system(f"copy {source_path} {destination_path}")
+        destination_path = results_path.joinpath(f"{graph_db_name}/config.yml")
+        shutil.copy2(source_path, destination_path)
 
 
 def preprocess_graph_data(experiment_configuration:dict):
diff --git a/src/framework/model_configuration.py b/src/framework/model_configuration.py
index bf3adb0..3ba44a9 100644
--- a/src/framework/model_configuration.py
+++ b/src/framework/model_configuration.py
@@ -14,7 +14,7 @@
 from framework.utils.data_sampling import curriculum_sampling
 from framework.utils.parameters import Parameters
 from models import GraphModel
-from src.utils.utils import get_k_lowest_nonzero_indices, valid_pruning_configuration, is_pruning
+from utils.utils import get_k_lowest_nonzero_indices, valid_pruning_configuration, is_pruning
 from utils.timer import TimeClass
 
 
@@ -59,6 +59,8 @@ def __init__(self, run_id: int, k_val: int, graph_data: GraphDataset, model_data
         self.scheduler = None
         self.net = None
         self.class_weights = None
+        self._csv_buffer = []
+        self._csv_flush_interval = self.para.run_config.config.get('csv_flush_interval', 10)
         # get gpu or cpu: (cpu is recommended at the moment)
         if self.para.run_config.config.get('device', None) is not None:
             self.device = torch.device(self.para.run_config.config['device'] if torch.cuda.is_available() else "cpu")
@@ -102,6 +104,7 @@ def train_configuration(self, pretrained_network=None):
             # Test early stopping criterion
             if self.early_stopping(epoch):
                 print(f"Early stopping at epoch {epoch}")
+                self._flush_csv_buffer()
                 break
 
             timer.measure("epoch")
@@ -141,10 +144,12 @@ def train_configuration(self, pretrained_network=None):
                     self.scheduler.step()
 
             # Evaluate the results on training, validation and test set (only if specified in the config or for evaluation)
-            if self.validate_data.size != 0:
+            validation_frequency = self.para.run_config.config.get('validation_frequency', 1)
+            is_validation_epoch = (epoch + 1) % validation_frequency == 0 or epoch == self.para.n_epochs - 1
+            if is_validation_epoch and self.validate_data.size != 0:
                 epoch_values, validation_values, test_values = self.evaluate_results(epoch=epoch,train_values=epoch_values, validation_values=validation_values, test_values=test_values, evaluation_type='validation')
             # check wheter there is a test split
-            if self.test_data.size != 0:
+            if is_validation_epoch and self.test_data.size != 0:
                 epoch_values, validation_values, test_values = self.evaluate_results(epoch=epoch,train_values=epoch_values, validation_values=validation_values, test_values=test_values, evaluation_type='test')
 
             timer.measure("epoch")
@@ -775,11 +780,19 @@ def postprocess_writer(self, epoch, epoch_time, train_values: EvaluationValues,
                             f"{validation_values.accuracy};{validation_values.loss};" \
                             f"{test_values.accuracy};{test_values.loss}\n"
 
-        # Save file for results
+        # Buffer CSV writes and flush periodically
+        self._csv_buffer.append(res_str)
+        if len(self._csv_buffer) >= self._csv_flush_interval or epoch == self.para.n_epochs - 1:
+            self._flush_csv_buffer()
+
+    def _flush_csv_buffer(self):
+        if not self._csv_buffer:
+            return
         file_name = f'{self.para.db}_{self.para.config_id}_Results_run_id_{self.run_id}_validation_step_{self.para.validation_id}.csv'
         final_path = self.results_path.joinpath(f'{self.para.db}/Results/{file_name}')
         with open(final_path, "a") as file_obj:
-            file_obj.write(res_str)
+            file_obj.writelines(self._csv_buffer)
+        self._csv_buffer.clear()
 
 
     def train_graph_task(self, epoch, values, train_batches, timer):
@@ -856,9 +869,9 @@ def evaluate_graph_task(self, graph_ids):
 
         with torch.no_grad():
             self.net.train(False)
-            # Run ordinary GNN
-            # split the graph ids into batches of size 512 to avoid memory issues
-            batches = [graph_ids[i:i + 512] for i in range(0, len(graph_ids), 512)]
+            # split the graph ids into batches to avoid memory issues
+            eval_batch_size = self.para.run_config.config.get('eval_batch_size', 512)
+            batches = [graph_ids[i:i + eval_batch_size] for i in range(0, len(graph_ids), eval_batch_size)]
             # deduce output size from batches and graph data
             total_out_len = 0
             for batch in batches:
@@ -870,6 +883,7 @@ def evaluate_graph_task(self, graph_ids):
             loader = CustomBatchLoader(self.graph_data, batches)
             batch_counter = 0
             if not self.para.run_config.config.get('with_invariant_layers', True):
+                # Run ordinary GNN
                 for i, batch in enumerate(loader):
                     #print(f"Evaluating batch {i + 1}/{len(batches)}")
                     outputs[batch_counter:batch_counter + len(batch)] = self.net(batch_data=batch)
@@ -877,7 +891,6 @@ def evaluate_graph_task(self, graph_ids):
             else:
                 # Run Share GNN
                 for j, data_pos in enumerate(graph_ids):
-                    self.net.train(False)
                     outputs[j] = self.net(self.graph_data[data_pos], pos=data_pos)
         return labels, outputs
 
diff --git a/src/framework/utils/evaluation.py b/src/framework/utils/evaluation.py
index 32dfe58..b66a95d 100644
--- a/src/framework/utils/evaluation.py
+++ b/src/framework/utils/evaluation.py
@@ -4,7 +4,7 @@
 import pandas as pd
 from matplotlib import pyplot as plt
 
-from src.utils.utils import is_pruning
+from utils.utils import is_pruning
 
 
 def epoch_accuracy(db_name, y_val, ids):
diff --git a/src/models/ShareGNN/layers/inv_based_message_passing.py b/src/models/ShareGNN/layers/inv_based_message_passing.py
index 676bae7..1acbfa9 100644
--- a/src/models/ShareGNN/layers/inv_based_message_passing.py
+++ b/src/models/ShareGNN/layers/inv_based_message_passing.py
@@ -68,12 +68,11 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase
         # Determine the number of weights and biases
         # There are two cases asymetric and symmetric, asymetric is the default, TODO add symmetric case
         self.skips = [0]
+        self.b_head_offset = 0
         self.skips_description = [None]
         self.skips_description_text = [None]
-        weight_dist_default_size = 1000000
-        self.weight_distribution = torch.zeros((weight_dist_default_size, 4), dtype=torch.int64).to(self.device)
-        last_weight_distribution_index = 0
-        self.bias_distribution = [None] * len(graph_data)
+        weight_distribution_chunks = [[] for _ in range(len(graph_data))]
+        bias_distribution_chunks = [[] for _ in range(len(graph_data))]
 
         # Iterate over all heads in the layer
         for head_id, head in enumerate(self.layer.layer_heads):
@@ -83,6 +82,10 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase
             source_labels = self.graph_data.node_labels[self.source_label_descriptions[head_id]].node_labels
             target_labels = self.graph_data.node_labels[self.target_label_descriptions[head_id]].node_labels
             bias_labels = self.graph_data.node_labels[self.bias_label_descriptions[head_id]].node_labels
+            current_head_id = 0
+            for h_i in range(head_id):
+                current_head_id += self.n_heads_per_label[h_i]
+
             for key in valid_property_values:
                 #print(f'Initialize head {i+1}/{len(self.layer.layer_heads)} with property {key}')
                 property_subdict = self.graph_data.properties[self.property_descriptions[head_id]].properties[key]
@@ -119,56 +122,76 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase
                     # relabel indices
                     indices[valid_indices] = torch.tensor([valid_value_dict[idx.item()] for idx in indices[valid_indices]], dtype=torch.int64)
                     num_weights = len(valid_values)
+                for n in range(self.n_heads_per_label[head_id]):
+                    self.weight_num.append(num_weights)
                 start_time = time.time()
-                for j in range(head.num):
-                    for idx in range(len(graph_data)):
-                        # if number of graphs is larger than 10000 print progress
-                        if len(graph_data) > 10000 and idx % 1000 == 0:
-                            print(f'Head {head_id+1}/{len(self.layer.layer_heads)} with property {key}: {idx}/{len(graph_data)} graphs processed ({(idx/len(graph_data))*100:.2f}%) time so far (in s): {time.time()-start_time:.2f}', flush=True)
-                        # get the valid indices for the current graph
-                        if threshold > 1 or do_invalid_indices_exist or upper_threshold is not None:
-                            valid_indices_graph = torch.where(valid_indices_bool[property_subdict_slices[idx]:property_subdict_slices[idx+1]])[0] + property_subdict_slices[idx]
-                        else:
-                            valid_indices_graph = torch.arange(property_subdict_slices[idx], property_subdict_slices[idx+1], dtype=torch.int64)
-                        # create new tensor where each row is the concatenation of head_id, property_subdict_row, and indices
-
-                        # Use self.weight_distribution directly
-                        self.weight_distribution[last_weight_distribution_index:last_weight_distribution_index + len(valid_indices_graph), 0] = head_id
-                        self.weight_distribution[last_weight_distribution_index:last_weight_distribution_index + len(valid_indices_graph), 1:3] = property_subdict[valid_indices_graph] - self.graph_data.slices['x'][idx] # check if subtracting is necessary
-                        self.weight_distribution[last_weight_distribution_index:last_weight_distribution_index + len(valid_indices_graph), 3] = indices[valid_indices_graph] + self.skips[-1]
-
-                self.skips.append(self.skips[-1] + num_weights)
-                self.skips_description.append({'head:': head_id, 'property': key, 'weights': num_weights})
-                self.skips_description_text.append(f"Head {head_id} Property {key} has {num_weights} different weights")
-
-
-            self.weight_num.append(self.skips[-1])
+                for idx in range(len(graph_data)):
+                    # if number of graphs is larger than 10000 print progress
+                    if len(graph_data) > 5000 and idx % 1000 == 0:
+                        print(
+                            f'Heads {self.n_heads_per_label[head_id] + current_head_id}/{self.num_heads} with property {key}: {idx}/{len(graph_data)} graphs processed ({(idx / len(graph_data)) * 100:.2f}%) time so far (in s): {time.time() - start_time:.2f}',
+                            flush=True)
+                    # get the valid indices for the current graph
+                    if threshold > 1 or do_invalid_indices_exist or upper_threshold is not None:
+                        valid_indices_graph = torch.where(
+                            valid_indices_bool[property_subdict_slices[idx]:property_subdict_slices[idx + 1]])[0] + \
+                                              property_subdict_slices[idx]
+                    else:
+                        valid_indices_graph = torch.arange(property_subdict_slices[idx],
+                                                           property_subdict_slices[idx + 1], dtype=torch.int64)
+                    w_indices = indices[valid_indices_graph]
+                    p_indices = property_subdict[valid_indices_graph] - self.graph_data.slices['x'][idx]  # check if subtracting is necessary
+                    # create new tensor where each row is the concatenation of head_id, property_subdict_row, and indices
+                    for n in range(self.n_heads_per_label[head_id]):
+                        new_weight_distribution = torch.zeros((len(valid_indices_graph), 4), dtype=torch.int64)
+                        new_weight_distribution[:, 0] = current_head_id + n
+                        new_weight_distribution[:, 1:3] = p_indices
+                        new_weight_distribution[:, 3] = w_indices + self.skips[-1] + n*num_weights
+                        weight_distribution_chunks[idx].append(new_weight_distribution)
+
+                for n in range(self.n_heads_per_label[head_id]):
+                    self.skips.append(self.skips[-1] + num_weights)
+                    self.skips_description.append({'head:': head_id, 'property': key, 'weights': num_weights})
+                    self.skips_description_text.append(f"Head {head_id} Property {key} has {num_weights} different weights")
+
             # TODO symmetric case
 
             if self.bias:
-                # Determine the number of different learnable parameters in the bias vector
-                self.bias_num.append(self.in_features * self.n_bias_labels[head_id])
                 # Set the bias weights
                 _, indices, counts = torch.unique(bias_labels, dim=0, return_inverse=True, return_counts=True, sorted=False)
                 for idx in range(len(graph_data)):
-                    for feature_id in range(self.in_features):
-                        new_bias_distribution = torch.zeros((graph_data.num_nodes[idx].item(), 4), dtype=torch.int64)
-                        new_bias_distribution[:, 0] = head_id
-                        new_bias_distribution[:, 1] = torch.arange(graph_data.num_nodes[idx].item(), dtype=torch.int64) # alternative torch.arange(start=graph_data.slices['x'][idx], end=graph_data.slices['x'][idx+1], dtype=torch.int64)
-                        new_bias_distribution[:, 2] = feature_id
-                        new_bias_distribution[:, 3] = indices[graph_data.slices['x'][idx]:graph_data.slices['x'][idx+1]] + feature_id * self.n_bias_labels[head_id]
-                        if self.bias_distribution[idx] is None:
-                            self.bias_distribution[idx] = new_bias_distribution.detach().clone()
-                        else:
-                            self.bias_distribution[idx] = torch.cat((self.bias_distribution[idx], new_bias_distribution), dim=0)
+                    arranged_tensor = torch.arange(graph_data.num_nodes[idx].item(), dtype=torch.int64) # alternative torch.arange(start=graph_data.slices['x'][idx], end=graph_data.slices['x'][idx+1], dtype=torch.int64)
+                    w_indices = indices[graph_data.slices['x'][idx]:graph_data.slices['x'][idx+1]]
+                    for n in range(self.n_heads_per_label[head_id]):
+                        for feature_id in range(self.in_features):
+                            w_index_offset = n*self.in_features*self.n_bias_labels[head_id] + feature_id*self.n_bias_labels[head_id]
+                            new_bias_distribution = torch.zeros((graph_data.num_nodes[idx].item(), 4), dtype=torch.int64)
+                            new_bias_distribution[:, 0] = current_head_id + n
+                            new_bias_distribution[:, 1] = arranged_tensor
+                            new_bias_distribution[:, 2] = feature_id
+                            new_bias_distribution[:, 3] = w_indices + self.b_head_offset + w_index_offset
+                            bias_distribution_chunks[idx].append(new_bias_distribution)
+                # Determine the number of different learnable parameters in the bias vector
+                for n in range(self.n_heads_per_label[head_id]):
+                    self.bias_num.append(self.in_features * self.n_bias_labels[head_id])
+                    self.b_head_offset += self.bias_num[-1]
 
         # Merge the weight distribution of all graphs (creating additionally slicing information)
-        self.weight_distribution_slices = torch.tensor([0] + [len(w) for w in self.weight_distribution], dtype=torch.int64).cumsum(dim=0)
-        self.weight_distribution = torch.cat([self.weight_distribution[i] for i in range(len(graph_data))], dim=0).to(self.device)
+        # Single torch.cat per graph instead of repeated incremental concatenation
+        merged_weight_distributions = [
+            torch.cat(chunks, dim=0) if chunks else torch.zeros((0, 4), dtype=torch.int64)
+            for chunks in weight_distribution_chunks
+        ]
+        self.weight_distribution_slices = torch.tensor([0] + [len(w) for w in merged_weight_distributions], dtype=torch.int64).cumsum(dim=0)
+        self.weight_distribution = torch.cat(merged_weight_distributions, dim=0).to(self.device)
         if self.bias:
             # Merge the bias distribution of all graphs (creating additionally slicing information)
-            self.bias_distribution_slices = torch.tensor([0] + [len(b) for b in self.bias_distribution], dtype=torch.int64).cumsum(dim=0)
-            self.bias_distribution = torch.cat([self.bias_distribution[i] for i in range(len(graph_data))], dim=0).to(self.device)
+            merged_bias_distributions = [
+                torch.cat(chunks, dim=0) if chunks else torch.zeros((0, 4), dtype=torch.int64)
+                for chunks in bias_distribution_chunks
+            ]
+            self.bias_distribution_slices = torch.tensor([0] + [len(b) for b in merged_bias_distributions], dtype=torch.int64).cumsum(dim=0)
+            self.bias_distribution = torch.cat(merged_bias_distributions, dim=0).to(self.device)
 
 
         if self.bias:
@@ -187,6 +210,10 @@ def __init__(self, parameters: Parameters, layer: Layer, graph_data: GraphDatase
 
         self.forward_step_time = 0
 
+        # Cache config lookups used in forward pass
+        self.use_degree_matrix = self.para.run_config.config.get('degree_matrix', False)
+        self.use_in_degrees = self.para.run_config.config.get('use_in_degrees', False)
+
     def init_weights(self, num_weights:np.float64, init_type:Optional[str]=None) -> nn.Parameter:
         """
         Initializes the weights, i.e., learnable parameters of the module
@@ -236,7 +263,7 @@ def set_weights(self, pos:int) -> None:
         :return:
         """
         input_size = self.graph_data.num_nodes[pos].item()
-        self.current_W = torch.zeros((self.num_heads, input_size, input_size), dtype=self.precision).to(self.device)
+        self.current_W = torch.zeros((self.num_heads, input_size, input_size), dtype=self.precision, device=self.device)
         graph_weight_distribution = self.weight_distribution[self.weight_distribution_slices[pos]:self.weight_distribution_slices[pos+1]]
         if len(graph_weight_distribution) != 0:
             # get third column of the weight_distribution: the index of self.Param_W
@@ -253,7 +280,7 @@ def set_bias(self, pos) -> None:
         :return:
         """
         input_size = self.graph_data.num_nodes[pos].item()
-        self.current_B = torch.zeros((self.num_heads, input_size, self.in_features), dtype=self.precision).to(self.device)
+        self.current_B = torch.zeros((self.num_heads, input_size, self.in_features), dtype=self.precision, device=self.device)
         graph_bias_distribution = self.bias_distribution[self.bias_distribution_slices[pos]:self.bias_distribution_slices[pos+1]]
         param_indices = graph_bias_distribution[:, 3]
         matrix_indices = graph_bias_distribution[:, 0:3].T
@@ -303,9 +330,9 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a
         begin = time.time()
         # set the weights, i.e., sets self.current_W to (C, N, N) where C is the number of channels and N is the number of nodes in graph at position pos of the dataset
         self.set_weights(pos)
-        if self.para.run_config.config.get('degree_matrix', False):
+        if self.use_degree_matrix:
             node_representation = self.in_edges[pos]*torch.einsum('cij,jk->cik', torch.diag(self.D[pos]) @ self.current_W @ torch.diag(self.D[pos]), node_representation)
-        elif self.para.run_config.config.get('use_in_degrees', False):
+        elif self.use_in_degrees:
             node_representation = self.in_edges[pos]*torch.einsum('cij,jk->cik', self.current_W, node_representation)
         else:
             node_representation = torch.einsum('hij,jf->hif', self.current_W, node_representation)
diff --git a/src/models/ShareGNN/layers/share_gnn_linear.py b/src/models/ShareGNN/layers/share_gnn_linear.py
index 7f674a3..11500cd 100644
--- a/src/models/ShareGNN/layers/share_gnn_linear.py
+++ b/src/models/ShareGNN/layers/share_gnn_linear.py
@@ -30,9 +30,6 @@ def __init__(self, seed:int, layer_id:int, layer:Layer, parameters, graph_data:G
             self.in_features = in_features
         self.in_features = layer.layer_dict.get('in_features', self.in_features)
 
-        self.out_features = graph_data.num_node_features
-        if out_features is not None:
-            self.out_features = out_features
         self.out_features = layer.layer_dict.get('out_features', self.in_features)
 
         # determine the number of heads
diff --git a/src/models/ShareGNN/preprocessing/preprocessing.py b/src/models/ShareGNN/preprocessing/preprocessing.py
index 7f73ede..6b7f367 100644
--- a/src/models/ShareGNN/preprocessing/preprocessing.py
+++ b/src/models/ShareGNN/preprocessing/preprocessing.py
@@ -150,7 +150,8 @@ def layer_to_labels(experiment_configuration, layer_strings: json, graph_data: G
                     raise ValueError(
                         f'Please specify the subgraphs in the config files under the key "subgraphs" as folllows: subgraphs: - "[nx.complete_graph(4)]"')
                 else:
-                    subgraph_list = eval(experiment_configuration['subgraphs'][layer['id']])
+                    import ast
+                    subgraph_list = ast.literal_eval(experiment_configuration['subgraphs'][layer['id']])
                     file_path = save_subgraph_labels(graph_data=graph_data,
                                                      subgraphs=subgraph_list,
                                                      subgraph_id=layer['id'],
diff --git a/src/models/ShareGNN/preprocessing/properties.py b/src/models/ShareGNN/preprocessing/properties.py
index 8973101..c35e7d8 100644
--- a/src/models/ShareGNN/preprocessing/properties.py
+++ b/src/models/ShareGNN/preprocessing/properties.py
@@ -11,7 +11,7 @@
 
 from datasets.graph_dataset import GraphDataset
 from datasets.utils.node_labeling import load_labels
-from src.utils.utils import convert_to_list
+from utils.utils import convert_to_list
 
 
 def write_distance_properties(graph_data:GraphDataset, cutoff=None, out_path: Path = Path(), save_times=None) -> None:
diff --git a/src/models/layers/framework_layer.py b/src/models/layers/framework_layer.py
index 5ccc9ec..02fba87 100644
--- a/src/models/layers/framework_layer.py
+++ b/src/models/layers/framework_layer.py
@@ -92,7 +92,28 @@ def __init__(self, layer_args=None):
         if 'activation' not in layer_args:
             self.activation = torch.nn.Identity()
         else:
-            self.activation = eval(layer_args['activation'])
+            activation_value = layer_args['activation']
+            if isinstance(activation_value, torch.nn.Module):
+                self.activation = activation_value
+            elif isinstance(activation_value, str):
+                ACTIVATION_MAP = {
+                    'torch.nn.Identity()': torch.nn.Identity(),
+                    'torch.nn.ReLU()': torch.nn.ReLU(),
+                    'torch.nn.LeakyReLU()': torch.nn.LeakyReLU(),
+                    'torch.nn.ELU()': torch.nn.ELU(),
+                    'torch.nn.GELU()': torch.nn.GELU(),
+                    'torch.nn.Tanh()': torch.nn.Tanh(),
+                    'torch.nn.Sigmoid()': torch.nn.Sigmoid(),
+                    'torch.nn.Softmax(dim=1)': torch.nn.Softmax(dim=1),
+                    'torch.nn.SELU()': torch.nn.SELU(),
+                }
+                if activation_value in ACTIVATION_MAP:
+                    self.activation = ACTIVATION_MAP[activation_value]
+                else:
+                    raise ValueError(f"Unsupported activation function: '{activation_value}'. "
+                                     f"Supported values: {list(ACTIVATION_MAP.keys())}")
+            else:
+                raise ValueError(f"Activation must be a string or torch.nn.Module, got {type(activation_value)}")
 
     @abstractmethod
     def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *args, **kwargs):
diff --git a/src/models/layers/mpnn_classical/gat_conv.py b/src/models/layers/mpnn_classical/gat_conv.py
index 87adb0c..4b90ef5 100644
--- a/src/models/layers/mpnn_classical/gat_conv.py
+++ b/src/models/layers/mpnn_classical/gat_conv.py
@@ -41,5 +41,5 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a
             else: # add residual to each head separately
                 node_representation = node_representation + x.repeat(1, self.gat_args['heads'])
         if self.dropout > 0 and self.training:
-            node_representation = torch.nn.Dropout(self.dropout)(node_representation)
+            node_representation = torch.nn.functional.dropout(node_representation, p=self.dropout, training=self.training)
         return node_representation
\ No newline at end of file
diff --git a/src/models/layers/mpnn_classical/gatv2_conv.py b/src/models/layers/mpnn_classical/gatv2_conv.py
index 89d66a5..786ae0d 100644
--- a/src/models/layers/mpnn_classical/gatv2_conv.py
+++ b/src/models/layers/mpnn_classical/gatv2_conv.py
@@ -46,5 +46,5 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a
             else: # add residual to each head separately
                 node_representation = node_representation + x.repeat(1, self.gatv2_args['heads'])
         if self.dropout > 0 and self.training:
-            node_representation = torch.nn.Dropout(self.dropout)(node_representation)
+            node_representation = torch.nn.functional.dropout(node_representation, p=self.dropout, training=self.training)
         return node_representation
\ No newline at end of file
diff --git a/src/models/layers/mpnn_classical/gcn_conv.py b/src/models/layers/mpnn_classical/gcn_conv.py
index f12e4b4..5b515b3 100644
--- a/src/models/layers/mpnn_classical/gcn_conv.py
+++ b/src/models/layers/mpnn_classical/gcn_conv.py
@@ -28,5 +28,5 @@ def forward(self,node_representation:torch.Tensor, batch_data: GraphDataset, *ar
         if self.residual:
             node_representation = node_representation + x
         if self.dropout > 0 and self.training:
-            node_representation = torch.nn.Dropout(self.dropout)(node_representation)
+            node_representation = torch.nn.functional.dropout(node_representation, p=self.dropout, training=self.training)
         return node_representation
\ No newline at end of file
diff --git a/src/models/layers/mpnn_classical/gin_conv.py b/src/models/layers/mpnn_classical/gin_conv.py
index 2d65c7c..a096079 100644
--- a/src/models/layers/mpnn_classical/gin_conv.py
+++ b/src/models/layers/mpnn_classical/gin_conv.py
@@ -38,5 +38,5 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a
         if self.residual:
             node_representation = node_representation + x
         if self.dropout > 0 and self.training:
-            node_representation = torch.nn.Dropout(self.dropout)(node_representation)
+            node_representation = torch.nn.functional.dropout(node_representation, p=self.dropout, training=self.training)
         return node_representation
\ No newline at end of file
diff --git a/src/models/layers/mpnn_classical/sage_conv.py b/src/models/layers/mpnn_classical/sage_conv.py
index f04bab1..8e84a94 100644
--- a/src/models/layers/mpnn_classical/sage_conv.py
+++ b/src/models/layers/mpnn_classical/sage_conv.py
@@ -29,5 +29,5 @@ def forward(self, node_representation:torch.Tensor, batch_data: GraphDataset, *a
         if self.residual:
             node_representation = node_representation + x
         if self.dropout > 0 and self.training:
-            node_representation = torch.nn.Dropout(self.dropout)(node_representation)
+            node_representation = torch.nn.functional.dropout(node_representation, p=self.dropout, training=self.training)
         return node_representation
\ No newline at end of file
diff --git a/src/models/model.py b/src/models/model.py
index 600f28a..c095d9c 100644
--- a/src/models/model.py
+++ b/src/models/model.py
@@ -69,14 +69,8 @@ def forward(self, batch_data, *args, **kwargs):
                                                 dtype=torch.double)
             x = x + random_variation
 
-        representation_list = []
         for i, layer in enumerate(self.net_layers):
             x = layer(x, batch_data, *args, **kwargs)
-            continue
-            #if isinstance(layer, GNNConvLayer):
-            #    representation_list.append(x)
-            #    if i == len(self.net_layers) - 1 or not isinstance(self.net_layers[i + 1], GNNConvLayer):
-            #        x = torch.squeeze(torch.mean(torch.stack(representation_list), 0, True), 0)
 
         return x
 
@@ -130,16 +124,17 @@ def get_model_layer(self, layer):
             layer_args = layer.layer_dict
             # GNN specific layers
             if layer.layer_type == LayerTypes.GCN_CONVOLUTION.value:
-                self.net_layers.append(GCNConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad))
+                return GCNConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad)
             elif layer.layer_type == LayerTypes.GAT_CONVOLUTION.value:
-                self.net_layers.append(GATConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad))
+                return GATConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad)
             elif layer.layer_type == LayerTypes.GATv2_CONVOLUTION.value:
-                self.net_layers.append(GATv2Conv(layer_args).type(self.precision).requires_grad_(self.convolution_grad))
+                return GATv2Conv(layer_args).type(self.precision).requires_grad_(self.convolution_grad)
             elif layer.layer_type == LayerTypes.GIN_CONVOLUTION.value:
-                self.net_layers.append(GINConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad))
+                return GINConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad)
             elif layer.layer_type == LayerTypes.SAGE_CONVOLUTION.value:
                 return SAGEConv(layer_args).type(self.precision).requires_grad_(self.convolution_grad)
 
+
         elif layer.layer_type == LayerTypes.GLOBAL_POOLING.value:
             layer_args = {'mode': layer.layer_dict.get('mode', 'mean')}
             return GlobalPooling(layer_args).type(self.precision).requires_grad_(self.aggregation_grad)