diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
new file mode 100644
index 0000000..0eaeb41
--- /dev/null
+++ b/.github/workflows/docs.yml
@@ -0,0 +1,47 @@
+name: docs
+
+on:
+  push:
+    branches: [main, polish-ecal-package]
+  pull_request:
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: pip
+
+      - name: Install package with docs extras
+        run: pip install -e ".[docs]"
+
+      - name: Build Sphinx docs
+        run: sphinx-build -b html -W --keep-going docs docs/_build/html
+
+      - name: Upload docs artifact
+        uses: actions/upload-pages-artifact@v3
+        with:
+          path: docs/_build/html
+
+  deploy:
+    if: github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && (github.ref == 'refs/heads/main' || github.ref == 'refs/heads/polish-ecal-package'))
+    needs: build
+    runs-on: ubuntu-latest
+    permissions:
+      pages: write
+      id-token: write
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@v4
diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 0000000..f24890a
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,2 @@
+recursive-include tests *.py
+recursive-include src/ecal/hardware/data *.yaml
diff --git a/README.md b/README.md
index 5f129a9..d9dbad5 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,8 @@ eCAL computes the total energy consumed across the full AI model lifecycle — d
 
 Published in **IEEE Journal on Selected Areas in Communications (JSAC), 2026**.
 
+[Documentation](https://vh2001.github.io/eCAL/)
+
 ## Installation
 
 ### From source (recommended for development)
@@ -84,7 +86,7 @@ python RunCalculator.py
 | Profile            | FP32 FLOPS   | TDP (W) | Device |
 |--------------------|-------------|---------|--------|
 | `apple_m2`         | 3.6 TFLOPS  | 22      | mps    |
-| `nvidia_a100_80gb` | 19.5 TFLOPS | 300     | cuda   |
+| `nvidia_a100_80gb` | 19.5 TFLOPS | 400     | cuda   |
 | `nvidia_h100_sxm`  | 67 TFLOPS   | 700     | cuda   |
 | `generic_cpu`      | 1 TFLOPS    | 100     | cpu    |
 | `generic_edge`     | 0.01 TFLOPS | 15      | cpu    |
diff --git a/calculators/ToyModels.py b/calculators/ToyModels.py
index ab2f27e..d75fbd4 100644
--- a/calculators/ToyModels.py
+++ b/calculators/ToyModels.py
@@ -1,24 +1,331 @@
-"""Backward-compatibility shim — use ecal.calculators.toy_models instead."""
-from ecal.calculators.toy_models import (  # noqa: F401
-    SimpleMLP,
-    SimpleCNN,
-    SimpleMLP_practical,
-    SimpleCNN_practical,
-    KANLikeRegressor,
-    MultiHeadSelfAttention,
-    TransformerBlock,
-    LossFun,
-    Transformer,
-)
-
-__all__ = [
-    "SimpleMLP",
-    "SimpleCNN",
-    "SimpleMLP_practical",
-    "SimpleCNN_practical",
-    "KANLikeRegressor",
-    "MultiHeadSelfAttention",
-    "TransformerBlock",
-    "LossFun",
-    "Transformer",
-]
+import torch.nn as nn
+import torch.nn.functional as F
+from efficient_kan import KAN
+
+import torch
+import numpy as np
+
+class SimpleMLP(nn.Module):
+    def __init__(self, input_size=10, hidden_size=10, output_size=2, num_layers=3):
+        super(SimpleMLP, self).__init__()
+
+        # Create a ModuleList to store variable number of layers
+        self.layers = nn.ModuleList()
+
+        # First layer (input to hidden)
+        self.layers.append(nn.Linear(input_size, hidden_size))
+        self.layers.append(nn.ReLU())
+
+        # Hidden layers
+        for _ in range(num_layers - 1):
+            self.layers.append(nn.Linear(hidden_size, hidden_size))
+            self.layers.append(nn.ReLU())
+
+        # Output layer
+        self.output = nn.Linear(hidden_size, output_size)
+
+    def forward(self, x):
+        # Pass through all layers sequentially
+        for layer in self.layers:
+            x = layer(x)
+        return self.output(x)
+class SimpleCNN(nn.Module):
+    def __init__(self, input_channels=1, hidden_channels=10, output_size=2, num_layers=3):
+        super(SimpleCNN, self).__init__()
+
+        self.layers = nn.ModuleList()
+        current_channels = input_channels
+
+        # Create convolutional layers
+        for i in range(num_layers):
+            # More controlled channel growth
+            out_channels = hidden_channels * (2 if i > 0 else 1)
+
+            # Create conv block with standard pooling
+            conv_block = nn.Sequential(
+                nn.Conv1d(current_channels, out_channels, kernel_size=3, padding=1),
+                nn.BatchNorm1d(out_channels),
+                nn.ReLU(),
+            )
+
+            self.layers.append(conv_block)
+            current_channels = out_channels
+
+        # Global average pooling
+        self.global_pool = nn.AvgPool1d(kernel_size=2)
+
+        # Output layer
+        self.output = nn.Linear(current_channels, output_size)
+
+    def forward(self, x):
+        # Pass through all convolutional blocks
+        for layer in self.layers:
+            x = layer(x)
+
+
+        # Ensure we have at least one feature
+        if x.size(-1) > 1:
+            x = self.global_pool(x)
+
+        # Global average pooling
+        x = torch.mean(x, dim=-1)
+
+        return self.output(x)
+
+
+class SimpleMLP_practical(nn.Module):
+    def __init__(self, input_size=10, hidden_size=10, output_size=2, num_layers=3):
+        super(SimpleMLP_practical, self).__init__()
+        self.layers = nn.ModuleList()
+        # Input layer
+        self.layers.append(nn.Linear(input_size, hidden_size))
+        self.layers.append(nn.ReLU())
+        # Hidden layers
+        for _ in range(num_layers - 2): # Adjusted loop for clarity
+            self.layers.append(nn.Linear(hidden_size, hidden_size))
+            self.layers.append(nn.ReLU())
+        # Output layer - Add it to the list!
+        self.layers.append(nn.Linear(hidden_size, output_size))
+
+    def forward(self, x):
+        # Simpler forward pass that processes all layers sequentially
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+# --- CNN Model (Modified to vary dense layers) ---
+class SimpleCNN_practical(nn.Module):
+    """
+    A simple 1D Convolutional Neural Network (CNN) for sequence or feature vector regression.
+    """
+    def __init__(self, input_size=10, output_size=2, hidden_channels=16, num_layers=3):
+        super(SimpleCNN_practical, self).__init__()
+        self.layers = nn.ModuleList()
+        # Input data is expected to be reshaped to have 1 channel.
+        current_channels = 1
+
+        # Create convolutional blocks with increasing channel depth
+        for i in range(1, num_layers):
+            out_channels = hidden_channels * (2 ** i)
+            conv_block = nn.Sequential(
+                nn.Conv1d(current_channels, out_channels, kernel_size=3, padding=1),
+                #nn.BatchNorm1d(out_channels),
+                #nn.ReLU()
+            )
+            self.layers.append(conv_block)
+            current_channels = out_channels
+
+        # Global pooling layer adapts to any input size
+        self.global_pool = nn.AdaptiveAvgPool1d(1)
+        self.output_layer = nn.Linear(current_channels, output_size)
+
+    def forward(self, x):
+
+        #if x.dim() == 2:
+        #    x = x.unsqueeze(1)  # Add channel dimension if missing
+        # Expected input x shape: (batch_size, 1, num_features)
+        for layer in self.layers:
+            x = layer(x)
+        #x = self.global_pool(x)
+        #x = x.view(x.size(0), -1) # Flatten the output for the linear layer
+        return x #self.output_layer(x)
+
+
+# --- Simplified KAN-like Model (Modified to vary sub-layers) ---
+class KANLikeRegressor(nn.Module):
+    def __init__(self, num_layers: int, grid_size: int = 10, din: int = 10, dout: int = 2, k: int = 3):
+        super(KANLikeRegressor, self).__init__()
+        self.input_dim = din
+
+        architecture = [din] + [din] * (num_layers) + [dout]
+
+        self.kan = KAN(architecture, grid_size=grid_size,
+                        spline_order=3)
+
+    def forward(self, x):
+        # Simpler forward pass that processes all layers sequentially
+        for layer in self.kan.layers:
+            x = layer(x)
+        return x
+
+
+#############################TRANSFORMER MODEL FROM SCRATCH#############################
+
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(self, num_emb, num_heads=8):
+        super().__init__()
+
+        # hyperparams
+        self.D = num_emb  # embedding size
+        self.H = num_heads # number of transformer heads
+
+        # weights for self-attention
+        self.w_k = nn.Linear(self.D, self.D * self.H)
+        self.w_q = nn.Linear(self.D, self.D * self.H)
+        self.w_v = nn.Linear(self.D, self.D * self.H)
+
+        # weights for a combination of multiple heads
+        self.w_c = nn.Linear(self.D * self.H, self.D)
+
+    def forward(self, x, causal=True):
+        # x: B(atch) x T(okens) x D(imensionality)
+        B, T, D = x.size()
+
+        # keys, queries, values
+        k = self.w_k(x).view(B, T, self.H, D) # B x T x H x D ########## K = x*W_k + b_k
+        q = self.w_q(x).view(B, T, self.H, D) # B x T x H x D ########## Q = x*W_q + b_q
+        v = self.w_v(x).view(B, T, self.H, D) # B x T x H x D ########## V = x*W_v + b_v
+
+        # batches and heads are merged for more efficent matrix multiplication
+        # B x T x H x D -> B*H x T x D
+        k = k.transpose(1, 2).contiguous().view(B * self.H, T, D) # B*H x T x D
+        q = q.transpose(1, 2).contiguous().view(B * self.H, T, D) # B*H x T x D
+        v = v.transpose(1, 2).contiguous().view(B * self.H, T, D) # B*H x T x D
+
+        k = k / (D**0.25) # scaling with sqrt(D)
+        q = q / (D**0.25) # scaling with sqrt(D)
+
+        # kq
+        kq = torch.bmm(q, k.transpose(1, 2)) # B*H x T x T # (Q x K^T) / sqrt(D)
+
+        # if causal apply mask to prevent information flow from future tokens we set tokens above the diagonal to -inf so after softmax they are 0
+        if causal:
+            mask = torch.triu_indices(T, T, offset=1)
+            kq[..., mask[0], mask[1]] = float('-inf')
+
+        # softmax
+        skq = F.softmax(kq, dim=2) # B*H x T x T | A = softmax((Q x K^T)/sqrt(D))
+
+        # self-attention
+        sa = torch.bmm(skq, v) # B*H x T x D # (softmax(Q x K^T) x V)
+        sa = sa.view(B, self.H, T, D) # B x H x T x D
+        sa = sa.transpose(1, 2) # B x T x H x D
+        sa = sa.contiguous().view(B, T, D * self.H) # B x T x D*H
+
+        out = self.w_c(sa) # B x T x D
+
+        return out
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, num_emb, num_neurons, num_heads=4):
+        super().__init__()
+
+        # hyperparams
+        self.D = num_emb
+        self.H = num_heads
+        self.neurons = num_neurons
+
+        # components
+        self.msha = MultiHeadSelfAttention(num_emb=self.D, num_heads=self.H)
+        self.layer_norm1 = nn.LayerNorm(self.D)
+        self.layer_norm2 = nn.LayerNorm(self.D)
+
+        self.mlp = nn.Sequential(nn.Linear(self.D, self.neurons * self.D),
+                                nn.GELU(),
+                                nn.Linear(self.neurons * self.D, self.D))
+
+    def forward(self, x, causal=True):
+        # Multi-Head Self-Attention
+        x_attn = self.msha(x, causal)
+        # LayerNorm
+        x = self.layer_norm1(x_attn + x)
+        # MLP
+        x_mlp = self.mlp(x)
+        # LayerNorm
+        x = self.layer_norm2(x_mlp + x)
+
+        return x
+
+
+class LossFun(nn.Module):
+    def __init__(self,):
+        super().__init__()
+
+        self.loss = nn.MSELoss()
+
+    def forward(self, y_model, y_true, reduction='sum'):
+        # y_model: B(atch) x T(okens) x V(alues)
+        # y_true: B x T
+        B, T, V = y_model.size()
+
+        y_model = y_model.view(B * T, V)
+        y_true = y_true.view(B * T,)
+
+
+        loss_matrix = self.loss(y_model, y_true) # B*T
+
+        if reduction == 'sum':
+            return torch.sum(loss_matrix)
+        elif reduction == 'mean':
+            loss_matrix = loss_matrix.view(B, T)
+            return torch.mean(torch.sum(loss_matrix, 1))
+        else:
+            raise ValueError('Reduction could be either `sum` or `mean`.')
+
+class Transformer(nn.Module):
+    def __init__(self, num_tokens, num_token_vals, num_emb, num_neurons, num_heads=2, dropout_prob=0.1, num_blocks=10, device='cpu'):
+        super().__init__()
+
+        # hyperparams
+        self.device = device
+        self.num_tokens = num_tokens
+        self.num_emb = num_emb
+        self.num_blocks = num_blocks
+
+        # FIX 2: Re-enable the positional embedding layer
+        self.positional_embedding = nn.Embedding(num_tokens, num_emb)
+
+        # This correctly projects each feature in the sequence to the embedding dimension
+        self.embedding = nn.Linear(1, num_emb)
+
+        # transformer blocks
+        self.transformer_blocks = nn.ModuleList()
+        for _ in range(num_blocks):
+            self.transformer_blocks.append(TransformerBlock(num_emb=num_emb, num_neurons=num_neurons, num_heads=num_heads))
+
+        # Output layer for regression (predicting 2 values: target_x, target_y)
+        regression_output_dim = 1
+        self.output_layer = nn.Sequential(nn.Linear(num_emb * num_tokens, regression_output_dim))
+
+        # dropout layer
+        self.dropout = nn.Dropout(dropout_prob)
+
+
+    def transformer_forward(self, x, causal=True):
+        # x starts as: B(atch) x T(okens) -> (32, 25)
+
+        # FIX 1: Add a dimension to the input for the linear embedding layer
+        # x becomes: B x T x 1 -> (32, 25, 1)
+        x = x.unsqueeze(-1)
+
+        # embedding of tokens
+        # x becomes B x T x D(im) -> (32, 25, 20)
+        x = self.embedding(x)
+
+        # embedding of positions
+        # Create positions tensor: (1, 25)
+        pos = torch.arange(0, x.shape[1], dtype=torch.long).unsqueeze(0).to(self.device)
+        # Get positional embeddings: (1, 25, 20)
+        pos_emb = self.positional_embedding(pos)
+
+        # Add positional embeddings to input embeddings
+        x = self.dropout(x + pos_emb)
+
+        # transformer blocks
+        for i in range(self.num_blocks):
+            x = self.transformer_blocks[i](x, causal=False) # Causal might not be needed for this task
+
+        # Flatten the output of the transformer blocks
+        # x shape from (B, T, D) to (B, T * D) -> (32, 25 * 20)
+        x = x.view(x.size(0), -1)
+
+        # output layer
+        out = self.output_layer(x)
+
+        return out
+
+    def forward(self, x, causal=True):
+        # This method just calls the main forward pass
+        return self.transformer_forward(x, causal=causal)
diff --git a/docs/configuration.md b/docs/configuration.md
index a528fff..11df9ea 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -20,7 +20,7 @@ used to convert FLOP estimates into energy. They are managed by
 | Profile            | FP32 FLOPS   | TDP (W) | Device |
 |--------------------|-------------|---------|--------|
 | `apple_m2`         | 3.6 TFLOPS  | 22      | mps    |
-| `nvidia_a100_80gb` | 19.5 TFLOPS | 300     | cuda   |
+| `nvidia_a100_80gb` | 19.5 TFLOPS | 400     | cuda   |
 | `nvidia_h100_sxm`  | 67 TFLOPS   | 700     | cuda   |
 | `generic_cpu`      | 1 TFLOPS    | 100     | cpu    |
 | `generic_edge`     | 0.01 TFLOPS | 15      | cpu    |
diff --git a/examples/basic_estimate.py b/examples/basic_estimate.py
new file mode 100644
index 0000000..8ba4eb7
--- /dev/null
+++ b/examples/basic_estimate.py
@@ -0,0 +1,31 @@
+"""Example: estimate the energy cost of an MLP's full AI lifecycle with eCAL.
+
+Run with:
+    python examples/basic_estimate.py
+"""
+
+import ecal
+
+# Model architecture: a 3-layer MLP with 10 input features and 2 output classes.
+model_params = {
+    "num_layers": 3,
+    "din": 10,
+    "dout": 2,
+}
+
+result = ecal.estimate(
+    model_type="MLP",
+    model_params=model_params,
+    num_samples=1000,          # training samples
+    sample_size=10,             # features per sample
+    num_epochs=50,
+    num_inferences=10000,       # inference calls to amortize the cost over
+    hardware="apple_m2",        # pick a profile from `ecal profiles`, or pass
+                                 # processor_flops_per_second / processor_max_power directly
+)
+
+print("Energy breakdown (Joules):")
+for stage in ("transmission", "preprocessing", "training", "evaluation", "inference_process"):
+    print(f"  {stage:18s}: {result[stage]:.6f} J")
+print(f"  {'total':18s}: {result['total']:.6f} J")
+print(f"\neCAL metric: {result['ecal_j_per_bit']:.3e} J/bit")
diff --git a/examples/full_pipeline.py b/examples/full_pipeline.py
new file mode 100644
index 0000000..b7cb2fd
--- /dev/null
+++ b/examples/full_pipeline.py
@@ -0,0 +1,124 @@
+"""Example: build an eCAL estimate from its individual pipeline stages.
+
+`ecal.estimate()` (see examples/basic_estimate.py) bundles all of this into
+one call. This script instead drives each stage directly through the
+underlying classes, which is useful when you need more control than
+`estimate()` exposes -- e.g. a custom multi-hop transmission path, or
+reusing one FLOP calculator across several estimates.
+
+Pipeline modeled here: an edge device collects samples, sends them over two
+network hops (WiFi -> Ethernet) to a server, which preprocesses, trains, and
+serves inference -- then ships inference results back over the same path.
+
+Run with:
+    python examples/full_pipeline.py
+"""
+
+from ecal.calculators.inference import Inference
+from ecal.calculators.model_flops import TransformerCalculator
+from ecal.calculators.preprocessing import DataPreprocessing
+from ecal.calculators.training import Training
+from ecal.calculators.transmission import Transmission
+from ecal.hardware.profiles import get_profile
+
+# --- Configuration -------------------------------------------------------
+
+hardware = get_profile("nvidia_h100_sxm")
+flops_per_second = hardware.flops_per_second_fp32
+processor_power = hardware.tdp_watts
+
+num_train_samples = 5000
+num_inferences = 20000
+sample_size = 32       # sequence length, used as the model's input dimension
+float_precision = 32    # bits per value transmitted/processed
+
+# A small decoder-only Transformer; reused for both training and inference
+# so its FLOP profile is only computed once per input shape.
+calculator = TransformerCalculator(
+    context_length=sample_size,
+    embedding_size=64,
+    num_heads=4,
+    num_decoder_blocks=4,
+    feed_forward_size=128,
+    vocab_size=5000,
+)
+input_size = (1, sample_size)
+
+# --- 1. Transmission: raw samples travel edge device -> gateway -> server,
+#        each hop with its own protocol stack. ---------------------------
+
+edge_to_gateway = Transmission(datalink="WIFI_MAC", physical="WIFI_PHY", failure_rate=0.01)
+gateway_to_server = Transmission(datalink="ETHERNET", physical="Generic_physical", failure_rate=0.0)
+
+train_bits = float_precision * sample_size * num_train_samples
+hop1 = edge_to_gateway.calculate_energy(train_bits)
+hop2 = gateway_to_server.calculate_energy(hop1["total_bits"])
+transmission_energy = hop1["total_energy"] + hop2["total_energy"]
+
+# --- 2. Preprocessing: normalize the data once it reaches the server. ----
+
+preprocessing = DataPreprocessing(
+    preprocessing_type="normalization",
+    processor_flops_per_second=flops_per_second,
+    processor_max_power=processor_power,
+)
+preprocessing_result = preprocessing.calculate_energy(num_train_samples, sample_size)
+preprocessing_energy = preprocessing_result["total_energy"]
+
+# --- 3. Training: 5-fold cross-validation over the preprocessed data. ----
+
+training = Training(
+    model_name="Transformer",
+    batch_size=32,
+    num_epochs=20,
+    num_samples=num_train_samples,
+    processor_flops_per_second=flops_per_second,
+    processor_max_power=processor_power,
+    input_size=input_size,
+    evaluation_strategy="cross_validation",
+    k_folds=5,
+    split_ratio=0.8,
+    calculator=calculator,
+)
+training_result = training.calculate_energy()
+
+# --- 4. Inference: serve num_inferences requests with the trained model. -
+
+inference = Inference(
+    model_name="Transformer",
+    input_size=input_size,
+    num_samples=num_inferences,
+    processor_flops_per_second=flops_per_second,
+    processor_max_power=processor_power,
+    calculator=calculator,
+)
+inference_energy = inference.calculate_energy()
+
+# --- 5. Ship inference results back over the same two-hop network. ------
+
+inference_bits = float_precision * sample_size * num_inferences
+inf_hop1 = edge_to_gateway.calculate_energy(inference_bits)
+inf_hop2 = gateway_to_server.calculate_energy(inf_hop1["total_bits"])
+inference_transmission_energy = inf_hop1["total_energy"] + inf_hop2["total_energy"]
+
+# --- Summary --------------------------------------------------------------
+
+total_energy = (
+    transmission_energy
+    + preprocessing_energy
+    + training_result["total_energy"]
+    + inference_energy
+    + inference_transmission_energy
+)
+total_bits = train_bits + inference_bits
+ecal_j_per_bit = total_energy / total_bits
+
+print("Energy breakdown (Joules):")
+print(f"  transmission (data upload) : {transmission_energy:.6f} J")
+print(f"  preprocessing              : {preprocessing_energy:.6f} J")
+print(f"  training                   : {training_result['training_energy']:.6f} J")
+print(f"  evaluation                 : {training_result['evaluation_energy']:.6f} J")
+print(f"  inference                  : {inference_energy:.6f} J")
+print(f"  transmission (results)     : {inference_transmission_energy:.6f} J")
+print(f"  {'total':27s}: {total_energy:.6f} J")
+print(f"\neCAL metric: {ecal_j_per_bit:.3e} J/bit")
diff --git a/pyproject.toml b/pyproject.toml
index a5c1fe0..edbc8bb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=68.0", "setuptools-scm>=8.0"]
+requires = ["setuptools>=68.0"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -10,7 +10,7 @@ readme = "README.md"
 license = "BSD-3-Clause"
 requires-python = ">=3.9"
 authors = [
-    {name = "SensorLab"},
+    {name = "SensorLab", email = "sensorlab.jsi@gmail.com"},
 ]
 keywords = ["energy", "AI", "lifecycle", "sustainability", "eCAL"]
 classifiers = [
@@ -56,7 +56,8 @@ docs = [
 ecal = "ecal.cli:main"
 
 [project.urls]
-Homepage = "https://github.com/cfortuna/eCAL"
+Homepage = "https://github.com/sensorlab/eCAL"
+Documentation = "https://vh2001.github.io/eCAL/"
 Paper = "https://ieeexplore.ieee.org/abstract/document/11298182"
 
 [tool.setuptools.packages.find]
diff --git a/src/ecal/calculators/toy_models.py b/src/ecal/calculators/toy_models.py
deleted file mode 100644
index d75fbd4..0000000
--- a/src/ecal/calculators/toy_models.py
+++ /dev/null
@@ -1,331 +0,0 @@
-import torch.nn as nn
-import torch.nn.functional as F
-from efficient_kan import KAN
-
-import torch
-import numpy as np
-
-class SimpleMLP(nn.Module):
-    def __init__(self, input_size=10, hidden_size=10, output_size=2, num_layers=3):
-        super(SimpleMLP, self).__init__()
-
-        # Create a ModuleList to store variable number of layers
-        self.layers = nn.ModuleList()
-
-        # First layer (input to hidden)
-        self.layers.append(nn.Linear(input_size, hidden_size))
-        self.layers.append(nn.ReLU())
-
-        # Hidden layers
-        for _ in range(num_layers - 1):
-            self.layers.append(nn.Linear(hidden_size, hidden_size))
-            self.layers.append(nn.ReLU())
-
-        # Output layer
-        self.output = nn.Linear(hidden_size, output_size)
-
-    def forward(self, x):
-        # Pass through all layers sequentially
-        for layer in self.layers:
-            x = layer(x)
-        return self.output(x)
-class SimpleCNN(nn.Module):
-    def __init__(self, input_channels=1, hidden_channels=10, output_size=2, num_layers=3):
-        super(SimpleCNN, self).__init__()
-
-        self.layers = nn.ModuleList()
-        current_channels = input_channels
-
-        # Create convolutional layers
-        for i in range(num_layers):
-            # More controlled channel growth
-            out_channels = hidden_channels * (2 if i > 0 else 1)
-
-            # Create conv block with standard pooling
-            conv_block = nn.Sequential(
-                nn.Conv1d(current_channels, out_channels, kernel_size=3, padding=1),
-                nn.BatchNorm1d(out_channels),
-                nn.ReLU(),
-            )
-
-            self.layers.append(conv_block)
-            current_channels = out_channels
-
-        # Global average pooling
-        self.global_pool = nn.AvgPool1d(kernel_size=2)
-
-        # Output layer
-        self.output = nn.Linear(current_channels, output_size)
-
-    def forward(self, x):
-        # Pass through all convolutional blocks
-        for layer in self.layers:
-            x = layer(x)
-
-
-        # Ensure we have at least one feature
-        if x.size(-1) > 1:
-            x = self.global_pool(x)
-
-        # Global average pooling
-        x = torch.mean(x, dim=-1)
-
-        return self.output(x)
-
-
-class SimpleMLP_practical(nn.Module):
-    def __init__(self, input_size=10, hidden_size=10, output_size=2, num_layers=3):
-        super(SimpleMLP_practical, self).__init__()
-        self.layers = nn.ModuleList()
-        # Input layer
-        self.layers.append(nn.Linear(input_size, hidden_size))
-        self.layers.append(nn.ReLU())
-        # Hidden layers
-        for _ in range(num_layers - 2): # Adjusted loop for clarity
-            self.layers.append(nn.Linear(hidden_size, hidden_size))
-            self.layers.append(nn.ReLU())
-        # Output layer - Add it to the list!
-        self.layers.append(nn.Linear(hidden_size, output_size))
-
-    def forward(self, x):
-        # Simpler forward pass that processes all layers sequentially
-        for layer in self.layers:
-            x = layer(x)
-        return x
-
-
-# --- CNN Model (Modified to vary dense layers) ---
-class SimpleCNN_practical(nn.Module):
-    """
-    A simple 1D Convolutional Neural Network (CNN) for sequence or feature vector regression.
-    """
-    def __init__(self, input_size=10, output_size=2, hidden_channels=16, num_layers=3):
-        super(SimpleCNN_practical, self).__init__()
-        self.layers = nn.ModuleList()
-        # Input data is expected to be reshaped to have 1 channel.
-        current_channels = 1
-
-        # Create convolutional blocks with increasing channel depth
-        for i in range(1, num_layers):
-            out_channels = hidden_channels * (2 ** i)
-            conv_block = nn.Sequential(
-                nn.Conv1d(current_channels, out_channels, kernel_size=3, padding=1),
-                #nn.BatchNorm1d(out_channels),
-                #nn.ReLU()
-            )
-            self.layers.append(conv_block)
-            current_channels = out_channels
-
-        # Global pooling layer adapts to any input size
-        self.global_pool = nn.AdaptiveAvgPool1d(1)
-        self.output_layer = nn.Linear(current_channels, output_size)
-
-    def forward(self, x):
-
-        #if x.dim() == 2:
-        #    x = x.unsqueeze(1)  # Add channel dimension if missing
-        # Expected input x shape: (batch_size, 1, num_features)
-        for layer in self.layers:
-            x = layer(x)
-        #x = self.global_pool(x)
-        #x = x.view(x.size(0), -1) # Flatten the output for the linear layer
-        return x #self.output_layer(x)
-
-
-# --- Simplified KAN-like Model (Modified to vary sub-layers) ---
-class KANLikeRegressor(nn.Module):
-    def __init__(self, num_layers: int, grid_size: int = 10, din: int = 10, dout: int = 2, k: int = 3):
-        super(KANLikeRegressor, self).__init__()
-        self.input_dim = din
-
-        architecture = [din] + [din] * (num_layers) + [dout]
-
-        self.kan = KAN(architecture, grid_size=grid_size,
-                        spline_order=3)
-
-    def forward(self, x):
-        # Simpler forward pass that processes all layers sequentially
-        for layer in self.kan.layers:
-            x = layer(x)
-        return x
-
-
-#############################TRANSFORMER MODEL FROM SCRATCH#############################
-
-class MultiHeadSelfAttention(nn.Module):
-    def __init__(self, num_emb, num_heads=8):
-        super().__init__()
-
-        # hyperparams
-        self.D = num_emb  # embedding size
-        self.H = num_heads # number of transformer heads
-
-        # weights for self-attention
-        self.w_k = nn.Linear(self.D, self.D * self.H)
-        self.w_q = nn.Linear(self.D, self.D * self.H)
-        self.w_v = nn.Linear(self.D, self.D * self.H)
-
-        # weights for a combination of multiple heads
-        self.w_c = nn.Linear(self.D * self.H, self.D)
-
-    def forward(self, x, causal=True):
-        # x: B(atch) x T(okens) x D(imensionality)
-        B, T, D = x.size()
-
-        # keys, queries, values
-        k = self.w_k(x).view(B, T, self.H, D) # B x T x H x D ########## K = x*W_k + b_k
-        q = self.w_q(x).view(B, T, self.H, D) # B x T x H x D ########## Q = x*W_q + b_q
-        v = self.w_v(x).view(B, T, self.H, D) # B x T x H x D ########## V = x*W_v + b_v
-
-        # batches and heads are merged for more efficent matrix multiplication
-        # B x T x H x D -> B*H x T x D
-        k = k.transpose(1, 2).contiguous().view(B * self.H, T, D) # B*H x T x D
-        q = q.transpose(1, 2).contiguous().view(B * self.H, T, D) # B*H x T x D
-        v = v.transpose(1, 2).contiguous().view(B * self.H, T, D) # B*H x T x D
-
-        k = k / (D**0.25) # scaling with sqrt(D)
-        q = q / (D**0.25) # scaling with sqrt(D)
-
-        # kq
-        kq = torch.bmm(q, k.transpose(1, 2)) # B*H x T x T # (Q x K^T) / sqrt(D)
-
-        # if causal apply mask to prevent information flow from future tokens we set tokens above the diagonal to -inf so after softmax they are 0
-        if causal:
-            mask = torch.triu_indices(T, T, offset=1)
-            kq[..., mask[0], mask[1]] = float('-inf')
-
-        # softmax
-        skq = F.softmax(kq, dim=2) # B*H x T x T | A = softmax((Q x K^T)/sqrt(D))
-
-        # self-attention
-        sa = torch.bmm(skq, v) # B*H x T x D # (softmax(Q x K^T) x V)
-        sa = sa.view(B, self.H, T, D) # B x H x T x D
-        sa = sa.transpose(1, 2) # B x T x H x D
-        sa = sa.contiguous().view(B, T, D * self.H) # B x T x D*H
-
-        out = self.w_c(sa) # B x T x D
-
-        return out
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, num_emb, num_neurons, num_heads=4):
-        super().__init__()
-
-        # hyperparams
-        self.D = num_emb
-        self.H = num_heads
-        self.neurons = num_neurons
-
-        # components
-        self.msha = MultiHeadSelfAttention(num_emb=self.D, num_heads=self.H)
-        self.layer_norm1 = nn.LayerNorm(self.D)
-        self.layer_norm2 = nn.LayerNorm(self.D)
-
-        self.mlp = nn.Sequential(nn.Linear(self.D, self.neurons * self.D),
-                                nn.GELU(),
-                                nn.Linear(self.neurons * self.D, self.D))
-
-    def forward(self, x, causal=True):
-        # Multi-Head Self-Attention
-        x_attn = self.msha(x, causal)
-        # LayerNorm
-        x = self.layer_norm1(x_attn + x)
-        # MLP
-        x_mlp = self.mlp(x)
-        # LayerNorm
-        x = self.layer_norm2(x_mlp + x)
-
-        return x
-
-
-class LossFun(nn.Module):
-    def __init__(self,):
-        super().__init__()
-
-        self.loss = nn.MSELoss()
-
-    def forward(self, y_model, y_true, reduction='sum'):
-        # y_model: B(atch) x T(okens) x V(alues)
-        # y_true: B x T
-        B, T, V = y_model.size()
-
-        y_model = y_model.view(B * T, V)
-        y_true = y_true.view(B * T,)
-
-
-        loss_matrix = self.loss(y_model, y_true) # B*T
-
-        if reduction == 'sum':
-            return torch.sum(loss_matrix)
-        elif reduction == 'mean':
-            loss_matrix = loss_matrix.view(B, T)
-            return torch.mean(torch.sum(loss_matrix, 1))
-        else:
-            raise ValueError('Reduction could be either `sum` or `mean`.')
-
-class Transformer(nn.Module):
-    def __init__(self, num_tokens, num_token_vals, num_emb, num_neurons, num_heads=2, dropout_prob=0.1, num_blocks=10, device='cpu'):
-        super().__init__()
-
-        # hyperparams
-        self.device = device
-        self.num_tokens = num_tokens
-        self.num_emb = num_emb
-        self.num_blocks = num_blocks
-
-        # FIX 2: Re-enable the positional embedding layer
-        self.positional_embedding = nn.Embedding(num_tokens, num_emb)
-
-        # This correctly projects each feature in the sequence to the embedding dimension
-        self.embedding = nn.Linear(1, num_emb)
-
-        # transformer blocks
-        self.transformer_blocks = nn.ModuleList()
-        for _ in range(num_blocks):
-            self.transformer_blocks.append(TransformerBlock(num_emb=num_emb, num_neurons=num_neurons, num_heads=num_heads))
-
-        # Output layer for regression (predicting 2 values: target_x, target_y)
-        regression_output_dim = 1
-        self.output_layer = nn.Sequential(nn.Linear(num_emb * num_tokens, regression_output_dim))
-
-        # dropout layer
-        self.dropout = nn.Dropout(dropout_prob)
-
-
-    def transformer_forward(self, x, causal=True):
-        # x starts as: B(atch) x T(okens) -> (32, 25)
-
-        # FIX 1: Add a dimension to the input for the linear embedding layer
-        # x becomes: B x T x 1 -> (32, 25, 1)
-        x = x.unsqueeze(-1)
-
-        # embedding of tokens
-        # x becomes B x T x D(im) -> (32, 25, 20)
-        x = self.embedding(x)
-
-        # embedding of positions
-        # Create positions tensor: (1, 25)
-        pos = torch.arange(0, x.shape[1], dtype=torch.long).unsqueeze(0).to(self.device)
-        # Get positional embeddings: (1, 25, 20)
-        pos_emb = self.positional_embedding(pos)
-
-        # Add positional embeddings to input embeddings
-        x = self.dropout(x + pos_emb)
-
-        # transformer blocks
-        for i in range(self.num_blocks):
-            x = self.transformer_blocks[i](x, causal=False) # Causal might not be needed for this task
-
-        # Flatten the output of the transformer blocks
-        # x shape from (B, T, D) to (B, T * D) -> (32, 25 * 20)
-        x = x.view(x.size(0), -1)
-
-        # output layer
-        out = self.output_layer(x)
-
-        return out
-
-    def forward(self, x, causal=True):
-        # This method just calls the main forward pass
-        return self.transformer_forward(x, causal=causal)
diff --git a/src/ecal/hardware/data/profiles.yaml b/src/ecal/hardware/data/profiles.yaml
index d0a7185..b9dae08 100644
--- a/src/ecal/hardware/data/profiles.yaml
+++ b/src/ecal/hardware/data/profiles.yaml
@@ -10,8 +10,8 @@ apple_m2:
 nvidia_a100_80gb:
   name: "NVIDIA A100 80GB"
   flops_per_second_fp32: 19.5e12
-  flops_per_second_fp16: 312e12
-  tdp_watts: 300
+  flops_per_second_fp16: 77.97e12
+  tdp_watts: 400
   gpu_power_watts: 250
   idle_power_watts: 50
   device: "cuda"
@@ -19,7 +19,7 @@ nvidia_a100_80gb:
 nvidia_h100_sxm:
   name: "NVIDIA H100 SXM"
   flops_per_second_fp32: 67e12
-  flops_per_second_fp16: 990e12
+  flops_per_second_fp16: 267.7e12
   tdp_watts: 700
   gpu_power_watts: 600
   idle_power_watts: 100