Repository: lhnguyen102/cuTAGI (language: python)

Issue:
Possible memory leakage ResNet CUDA
@lhnguyen102 Hi Ha, 

I am not sure if it is supposed to be like this, but I believe there is a memory leakage on the ResNet architecture in CUDA.

While running the resnet18_cifar10 example in the server, it runs out of memory on epoch number 4:
`Epoch 4/10 | training error: 43.58% | test error: 48.08%:  40%|██████████████▊                      | 4/10 [25:27<38:13, 382.20s/it]CUDA Runtime Error at: /home/mf/Documents/RESNET/cuTAGI/src/data_struct_cuda.cu:83
out of memory cudaMalloc(&this->d_jcb, size * sizeof(float))`

I tracked the memory usage during the epochs and this is what I got:
- Epoch 0/10: 3149MiB / 16384MiB
- Epoch 1/10: 8755MiB / 16384MiB
- Epoch 2/10: 10173MiB / 16384MiB
- Epoch 3/10: 13683MiB / 16384MiB


Discussion / hints:
Miquel, thanks for pointing out the issue. I will work on this issue this week
@miquelflorensa and @jamesgoulet, I confirmed the issue comes from python binding code I've tested C++ version and run memory check. There is no memory leak in C++/CUDA. I am working on the fix for Python binding codes

Fix patch:
diff --git a/.github/workflows/build_wheels.yml b/.github/workflows/build_wheels.yml
index 06c5c5d4d..559062b80 100644
--- a/.github/workflows/build_wheels.yml
+++ b/.github/workflows/build_wheels.yml
@@ -60,15 +60,15 @@ jobs:
     strategy:
       matrix:
         include:
-          - os: macos-11
+          - os: macos-13
             python_version: "3.9"
             cibuildwheel_python: "cp39-*"
 
-          - os: macos-11
+          - os: macos-13
             python_version: "3.10"
             cibuildwheel_python: "cp310-*"
 
-          - os: macos-11
+          - os: macos-13
             python_version: "3.11"
             cibuildwheel_python: "cp311-*"
     if: ${{ github.event.pull_request.draft == false }}
diff --git a/.gitignore b/.gitignore
index b924c162a..938bdebf0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,7 +11,7 @@ __pycache__/
 pytagi.egg-info
 dist
 data/cifar
-data/cifar_torch
+data/mnist/MNIST
 *ncu-rep
 *sys-rep
 *sqlite
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3512d855a..8e7eb0481 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -89,6 +89,7 @@ endif()
 # NOTE: This might need to change for higher CUDA version
 list(APPEND CUDA_NVCC_FLAGS "--expt-extended-lambda")
 list(APPEND CUDA_NVCC_FLAGS "--expt-relaxed-constexpr")
+# list(APPEND CUDA_NVCC_FLAGS "--use_fast_math")
 # list(APPEND CUDA_NVCC_FLAGS "-lineinfo")
 
 if(APPLE)
diff --git a/examples/autoencoder.py b/examples/autoencoder.py
index d594996c2..a34dece88 100644
--- a/examples/autoencoder.py
+++ b/examples/autoencoder.py
@@ -21,6 +21,7 @@
     Sequential,
 )
 
+
 plt.rcParams.update(
     {
         "font.size": 18,
diff --git a/examples/classification.py b/examples/classification.py
index 6cc7cbb83..da0a6516a 100644
--- a/examples/classification.py
+++ b/examples/classification.py
@@ -1,11 +1,3 @@
-# Temporary import. It will be removed in the final version
-import os
-import sys
-
-# Add the 'build' directory to sys.path in one line
-sys.path.append(
-    os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "build"))
-)
 import fire
 import numpy as np
 from tqdm import tqdm
@@ -116,9 +108,9 @@ def main(num_epochs: int = 10, batch_size: int = 256, sigma_v: float = 2.0):
     metric = HRCSoftmaxMetric(num_classes=10)
 
     # Network configuration
-    net = CNN_BATCHNORM
-    # net.to_device("cuda")
-    # net.set_threads(16)
+    net = CNN
+    net.to_device("cuda")
+    net.set_threads(16)
     out_updater = OutputUpdater(net.device)
 
     # Training
@@ -155,7 +147,7 @@ def main(num_epochs: int = 10, batch_size: int = 256, sigma_v: float = 2.0):
 
         # Testing
         test_error_rates = []
-        test_batch_iter = test_dtl.create_data_loader(40, shuffle=False)
+        test_batch_iter = test_dtl.create_data_loader(batch_size, shuffle=False)
         for x, _, _, label in test_batch_iter:
             m_pred, v_pred = net(x)
 
diff --git a/examples/mnist_bench.py b/examples/mnist_bench.py
index 7a5ceba51..b02326647 100644
--- a/examples/mnist_bench.py
+++ b/examples/mnist_bench.py
@@ -1,10 +1,13 @@
+import fire
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torchvision import datasets, transforms
 from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from tqdm import tqdm
 
-from pytagi import HRCSoftmaxMetric
+from pytagi import HRCSoftmaxMetric, Utils, exponential_scheduler
 from pytagi.nn import (
     AvgPool2d,
     BatchNorm2d,
@@ -16,83 +19,316 @@
     Sequential,
 )
 
-TAGINN = Sequential(
-    Linear(784, 100),
+TAGI_FNN = Sequential(
+    Linear(784, 4096),
     ReLU(),
-    Linear(100, 100),
+    Linear(4096, 4096),
     ReLU(),
-    Linear(100, 11),
+    Linear(4096, 11),
 )
 
+TAGI_CNN = Sequential(
+    Conv2d(1, 16, 4, padding=1, in_width=28, in_height=28),
+    ReLU(),
+    AvgPool2d(3, 2),
+    Conv2d(16, 32, 5),
+    ReLU(),
+    AvgPool2d(3, 2),
+    Linear(32 * 4 * 4, 256),
+    ReLU(),
+    Linear(256, 11),
+)
 
-# Define the neural network model
-class TorchNN(nn.Module):
+TAGI_CNN_BATCHNORM = Sequential(
+    Conv2d(1, 16, 4, padding=1, in_width=28, in_height=28, bias=False),
+    ReLU(),
+    BatchNorm2d(16),
+    AvgPool2d(3, 2),
+    Conv2d(16, 32, 5, bias=False),
+    ReLU(),
+    BatchNorm2d(32),
+    AvgPool2d(3, 2),
+    Linear(32 * 4 * 4, 256),
+    ReLU(),
+    Linear(256, 11),
+)
+
+
+DATA_FOLDER = "./data/mnist"
+
+
+# TORCH
+class TorchFNN(nn.Module):
     def __init__(self):
-        super(TorchNN, self).__init__()
-        self.fc1 = nn.Linear(28 * 28, 100)
-        self.fc2 = nn.Linear(100, 100)
-        self.fc3 = nn.Linear(100, 10)
+        super(TorchFNN, self).__init__()
+        self.model = nn.Sequential(
+            nn.Linear(28 * 28, 4096),
+            nn.ReLU(),
+            nn.Linear(4096, 4096),
+            nn.ReLU(),
+            nn.Linear(4096, 10),
+        )
 
     def forward(self, x):
         x = x.view(-1, 28 * 28)
-        x = torch.relu(self.fc1(x))
-        x = torch.relu(self.fc2(x))
-        x = self.fc3(x)
+        x = self.model(x)
         return x
 
 
-# Hyperparameters
-batch_size = 64
-learning_rate = 0.01
-num_epochs = 10
+class TorchCNN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.Conv2d(1, 16, kernel_size=4, padding=1),
+            nn.ReLU(),
+            nn.AvgPool2d(kernel_size=3, stride=2),
+            nn.Conv2d(16, 32, kernel_size=5),
+            nn.ReLU(),
+            nn.AvgPool2d(kernel_size=3, stride=2),
+            nn.Flatten(),
+            nn.Linear(32 * 4 * 4, 256),
+            nn.ReLU(),
+            nn.Linear(256, 10),
+        )
 
-# Data loading and preprocessing
-transform = transforms.Compose(
-    [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
-)
+    def forward(self, x):
+        return self.model(x)
 
-train_dataset = datasets.MNIST(
-    root="./data/torch", train=True, transform=transform, download=True
-)
-test_dataset = datasets.MNIST(
-    root="./data/torch", train=False, transform=transform, download=True
-)
 
-train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
-test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
-
-# Initialize the model, loss function, and optimizer
-model = TorchNN()
-criterion = nn.CrossEntropyLoss()
-optimizer = optim.SGD(model.parameters(), lr=learning_rate)
-
-# Training loop
-for epoch in range(num_epochs):
-    model.train()
-    for batch_idx, (data, target) in enumerate(train_loader):
-        optimizer.zero_grad()
-        output = model(data)
-        loss = criterion(output, target)
-        loss.backward()
-        optimizer.step()
-
-        if batch_idx % 100 == 0:
-            print(
-                f"Epoch: {epoch+1}/{num_epochs} [{batch_idx*len(data)}/{len(train_loader.dataset)}] Loss: {loss.item():.6f}"
+class TorchCNNBatchNorm(nn.Module):
+    def __init__(self):
+        super(TorchCNNBatchNorm, self).__init__()
+        self.model = nn.Sequential(
+            nn.Conv2d(1, 16, kernel_size=4, padding=1),
+            nn.BatchNorm2d(
+                16
+            ),  # Add batch normalization after the first convolutional layer
+            nn.ReLU(),
+            nn.AvgPool2d(kernel_size=3, stride=2),
+            nn.Conv2d(16, 32, kernel_size=5),
+            nn.BatchNorm2d(
+                32
+            ),  # Add batch normalization after the second convolutional layer
+            nn.ReLU(),
+            nn.AvgPool2d(kernel_size=3, stride=2),
+            nn.Flatten(),
+            nn.Linear(32 * 4 * 4, 256),
+            nn.ReLU(),
+            nn.Linear(256, 10),
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+
+def custom_collate_fn(batch):
+    # batch is a list of tuples (image, label)
+    batch_images, batch_labels = zip(*batch)
+
+    # Convert to a single tensor and then to numpy
+    batch_images = torch.stack(batch_images)
+    batch_labels = torch.tensor(batch_labels)
+
+    # Flatten images and labels to 1D
+    batch_images = batch_images.numpy().reshape(len(batch_images), -1).flatten()
+    batch_labels = batch_labels.numpy().flatten()
+
+    return batch_images, batch_labels
+
+
+def tagi_trainer(
+    batch_size: int, num_epochs: int, device: str = "cpu", sigma_v: float = 2.0
+):
+    # Data loading and preprocessing
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+    )
+
+    train_dataset = datasets.MNIST(
+        root=DATA_FOLDER, train=True, transform=transform, download=True
+    )
+    test_dataset = datasets.MNIST(
+        root=DATA_FOLDER, train=False, transform=transform, download=True
+    )
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=1,
+        collate_fn=custom_collate_fn,
+    )
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=1,
+        collate_fn=custom_collate_fn,
+    )
+
+    utils = Utils()
+
+    # Hierachical Softmax
+    metric = HRCSoftmaxMetric(num_classes=10)
+    net = TAGI_FNN
+    net.to_device(device)
+    # net.set_threads(16)
+    out_updater = OutputUpdater(net.device)
+
+    # Training
+    error_rates = []
+    pbar = tqdm(range(num_epochs), desc="Training Progress")
+
+    for epoch in pbar:
+        # count = 0
+
+        # Decaying observation's variance
+        sigma_v = exponential_scheduler(
+            curr_v=sigma_v, min_v=0.3, decaying_factor=0.99, curr_iter=epoch
+        )
+        var_y = np.full(
+            (batch_size * metric.hrc_softmax.num_obs,), sigma_v**2, dtype=np.float32
+        )
+        for x, labels in train_loader:
+            # Feedforward and backward pass
+            m_pred, v_pred = net(x)
+
+            # Update output layers based on targets
+            y, y_idx, _ = utils.label_to_obs(labels=labels, num_classes=10)
+            out_updater.update_using_indices(
+                output_states=net.output_z_buffer,
+                mu_obs=y,
+                var_obs=var_y,
+                selected_idx=y_idx,
+                delta_states=net.input_delta_z_buffer,
             )
 
-    # Evaluate the model on the test set
-    model.eval()
-    test_loss = 0
-    correct = 0
-    with torch.no_grad():
-        for data, target in test_loader:
+            # Update parameters
+            net.backward()
+            net.step()
+
+            # Training metric
+            error_rate = metric.error_rate(m_pred, v_pred, labels)
+            error_rates.append(error_rate)
+
+        # Averaged error
+        avg_error_rate = sum(error_rates[-100:])
+
+        # Testing
+        test_error_rates = []
+        for x, labels in test_loader:
+
+            m_pred, v_pred = net(x)
+
+            # Training metric
+            error_rate = metric.error_rate(m_pred, v_pred, labels)
+            test_error_rates.append(error_rate)
+
+        test_error_rate = sum(test_error_rates) / len(test_error_rates)
+        pbar.set_description(
+            f"Epoch {epoch + 1}/{num_epochs} | training error: {avg_error_rate:.2f}% | test error: {test_error_rate * 100:.2f}%",
+            refresh=True,
+        )
+    print("Training complete.")
+
+
+def torch_trainer(batch_size: int, num_epochs: int, device: str = "cpu"):
+    # Hyperparameters
+    learning_rate = 0.01
+
+    # torch.set_float32_matmul_precision("high")
+
+    # Data loading and preprocessing
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+    )
+
+    train_dataset = datasets.MNIST(
+        root=DATA_FOLDER, train=True, transform=transform, download=True
+    )
+    test_dataset = datasets.MNIST(
+        root=DATA_FOLDER, train=False, transform=transform, download=True
+    )
+
+    train_loader = DataLoader(
+        dataset=train_dataset, batch_size=batch_size, shuffle=True
+    )
+    test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
+
+    # Initialize the model, loss function, and optimizer
+    torch_device = torch.device(device)
+    if "cuda" in device and not torch.cuda.is_available():
+        raise RuntimeError(
+            "CUDA is not available. Please check your CUDA installation."
+        )
+    model = TorchFNN().to(torch_device)
+    # model = torch.compile(model, mode="reduce-overhead")
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(model.parameters(), lr=learning_rate)
+
+    if not torch.cuda.is_available() and "cuda" in device:
+        raise RuntimeError(
+            "CUDA is not available. Please check your CUDA installation."
+        )
+
+    # Training loop
+    pbar = tqdm(range(num_epochs), desc="Training Progress")
+    error_rates = []
+    for epoch in pbar:
+        model.train()
+        for _, (data, target) in enumerate(train_loader):
+            data = data.to(torch_device)
+            target = target.to(torch_device)
+
+            optimizer.zero_grad()
             output = model(data)
-            test_loss += criterion(output, target).item()
+            loss = criterion(output, target)
+
+            loss.backward()
+            optimizer.step()
+
             pred = output.argmax(dim=1, keepdim=True)
-            correct += pred.eq(target.view_as(pred)).sum().item()
+            train_correct = pred.eq(target.view_as(pred)).sum().item()
+            error_rates.append((1.0 - (train_correct / data.shape[0])))
 
-    test_loss /= len(test_loader.dataset)
-    print(
-        f"\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.0f}%)\n"
-    )
+        # Averaged error
+        avg_error_rate = sum(error_rates[-100:])
+
+        # Evaluate the model on the test set
+        model.eval()
+        test_loss = 0
+        correct = 0
+        with torch.no_grad():
+            for data, target in test_loader:
+                data = data.to(torch_device)
+                target = target.to(torch_device)
+                output = model(data)
+
+                test_loss += criterion(output, target).item()
+                pred = output.argmax(dim=1, keepdim=True)
+                correct += pred.eq(target.view_as(pred)).sum().item()
+
+        test_loss /= len(test_loader.dataset)
+        test_error_rate = (1.0 - correct / len(test_loader.dataset)) * 100
+        pbar.set_description(
+            f"Epoch# {epoch/num_epochs}| training error: {avg_error_rate:.2f}% | Test error: {test_error_rate: .2f}%"
+        )
+
+
+def main(
+    framework: str = "tagi",
+    batch_size: int = 256,
+    epochs: int = 20,
+    device: str = "cuda",
+):
+    if framework == "torch":
+        torch_trainer(batch_size=batch_size, num_epochs=epochs, device=device)
+    elif framework == "tagi":
+        tagi_trainer(batch_size=batch_size, num_epochs=epochs, device=device)
+    else:
+        raise RuntimeError(f"Invalid Framework: {framework}")
+
+
+if __name__ == "__main__":
+    fire.Fire(main)
diff --git a/examples/resnet18_cifar10.py b/examples/resnet18_cifar10.py
index 3b3bea929..72b6a2993 100644
--- a/examples/resnet18_cifar10.py
+++ b/examples/resnet18_cifar10.py
@@ -1,12 +1,3 @@
-# Temporary import. It will be removed in the final version
-import os
-import sys
-
-# Add the 'build' directory to sys.path in one line
-sys.path.append(
-    os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "build"))
-)
-
 import fire
 import numpy as np
 import torch
@@ -32,6 +23,23 @@
 NORMALIZATION_MEAN = (0.4914, 0.4822, 0.4465)
 NORMALIZATION_STD = (0.2023, 0.1994, 0.2010)
 
+CNN_NET = Sequential(
+    Conv2d(3, 32, 5, bias=False, padding=2, in_width=32, in_height=32),
+    BatchNorm2d(32),
+    ReLU(),
+    AvgPool2d(3, 2, padding=1, padding_type=2),
+    BatchNorm2d(32),
+    ReLU(),
+    AvgPool2d(3, 2, padding=1, padding_type=2),
+    Conv2d(32, 64, 5, bias=False, padding=2),
+    BatchNorm2d(64),
+    ReLU(),
+    AvgPool2d(3, 2, padding=1, padding_type=2),
+    Linear(64 * 4 * 4, 128),
+    ReLU(),
+    Linear(128, 11),
+)
+
 
 def make_layer_block(in_c: int, out_c: int, stride: int = 1, padding_type: int = 1):
     """Create a layer block for resnet 18"""
diff --git a/examples/time_series_forecasting.py b/examples/time_series_forecasting.py
index 139b2563d..c158ca633 100644
--- a/examples/time_series_forecasting.py
+++ b/examples/time_series_forecasting.py
@@ -1,12 +1,4 @@
-# Temporary import. It will be removed in the final vserion
-import os
-import sys
-
-# Add the 'build' directory to sys.path in one line
-sys.path.append(
-    os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "build"))
-)
-from typing import Union, Optional
+from typing import Optional
 
 import fire
 import matplotlib.pyplot as plt
diff --git a/include/activation.h b/include/activation.h
index 2b4a49cc7..1b4dcf044 100644
--- a/include/activation.h
+++ b/include/activation.h
@@ -127,18 +127,17 @@ class ReLU : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
+    using BaseLayer::backward;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
 #ifdef USE_CUDA
     std::unique_ptr<BaseLayer> to_cuda() override;
@@ -172,18 +171,17 @@ class Sigmoid : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
+    using BaseLayer::backward;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
 #ifdef USE_CUDA
     std::unique_ptr<BaseLayer> to_cuda() override;
@@ -217,18 +215,17 @@ class Tanh : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
+    using BaseLayer::backward;
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
 #ifdef USE_CUDA
     std::unique_ptr<BaseLayer> to_cuda() override;
@@ -262,18 +259,17 @@ class MixtureReLU : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
+    using BaseLayer::backward;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
 #ifdef USE_CUDA
     std::unique_ptr<BaseLayer> to_cuda() override;
@@ -307,18 +303,17 @@ class MixtureSigmoid : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
+    using BaseLayer::backward;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
 #ifdef USE_CUDA
     std::unique_ptr<BaseLayer> to_cuda() override;
@@ -352,18 +347,17 @@ class MixtureTanh : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
+    using BaseLayer::backward;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
 #ifdef USE_CUDA
     std::unique_ptr<BaseLayer> to_cuda() override;
@@ -397,18 +391,17 @@ class Softplus : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
+    using BaseLayer::backward;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
 #ifdef USE_CUDA
     std::unique_ptr<BaseLayer> to_cuda() override;
@@ -443,18 +436,17 @@ class LeakyReLU : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
+    using BaseLayer::backward;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
 #ifdef USE_CUDA
     std::unique_ptr<BaseLayer> to_cuda() override;
@@ -489,18 +481,17 @@ class Softmax : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
+    using BaseLayer::backward;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
 #ifdef USE_CUDA
     std::unique_ptr<BaseLayer> to_cuda() override;
@@ -544,18 +535,17 @@ class RemaxA : public BaseLayer {
 
     void forward(BaseHiddenStates &input_states,
                  BaseHiddenStates &output_states,
-                 BaseTempStates &temp_states) override{};
+                 BaseTempStates &temp_states) override {};
 
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
+    using BaseLayer::backward;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 };
\ No newline at end of file
diff --git a/include/activation_cuda.cuh b/include/activation_cuda.cuh
index 871d09a72..3d2137281 100644
--- a/include/activation_cuda.cuh
+++ b/include/activation_cuda.cuh
@@ -85,15 +85,17 @@ class ReLUCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void allocate_param_delta() override{};
+    using BaseLayer::backward;
 
-    void update_weights() override{};
+    void allocate_param_delta() override {};
 
-    void update_biases() override{};
+    void update_weights() override {};
 
-    void save(std::ofstream &file) override{};
+    void update_biases() override {};
 
-    void load(std::ifstream &file) override{};
+    void save(std::ofstream &file) override {};
+
+    void load(std::ifstream &file) override {};
 
     std::unique_ptr<BaseLayer> to_host() override;
 };
@@ -122,15 +124,17 @@ class SigmoidCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void allocate_param_delta() override{};
+    using BaseLayer::backward;
+
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
     std::unique_ptr<BaseLayer> to_host() override;
 };
@@ -160,15 +164,17 @@ class TanhCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void allocate_param_delta() override{};
+    using BaseLayer::backward;
+
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
     std::unique_ptr<BaseLayer> to_host() override;
 };
@@ -197,15 +203,17 @@ class MixtureReLUCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void allocate_param_delta() override{};
+    using BaseLayer::backward;
 
-    void update_weights() override{};
+    void allocate_param_delta() override {};
 
-    void update_biases() override{};
+    void update_weights() override {};
 
-    void save(std::ofstream &file) override{};
+    void update_biases() override {};
 
-    void load(std::ifstream &file) override{};
+    void save(std::ofstream &file) override {};
+
+    void load(std::ifstream &file) override {};
 
     std::unique_ptr<BaseLayer> to_host() override;
 };
@@ -234,15 +242,17 @@ class MixtureSigmoidCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void allocate_param_delta() override{};
+    using BaseLayer::backward;
+
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
     std::unique_ptr<BaseLayer> to_host() override;
 };
@@ -271,15 +281,17 @@ class MixtureTanhCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void allocate_param_delta() override{};
+    using BaseLayer::backward;
 
-    void update_weights() override{};
+    void allocate_param_delta() override {};
 
-    void update_biases() override{};
+    void update_weights() override {};
 
-    void save(std::ofstream &file) override{};
+    void update_biases() override {};
 
-    void load(std::ifstream &file) override{};
+    void save(std::ofstream &file) override {};
+
+    void load(std::ifstream &file) override {};
 
     std::unique_ptr<BaseLayer> to_host() override;
 };
@@ -308,15 +320,17 @@ class SoftplusCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void allocate_param_delta() override{};
+    using BaseLayer::backward;
+
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
     std::unique_ptr<BaseLayer> to_host() override;
 };
@@ -346,15 +360,17 @@ class LeakyReLUCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void allocate_param_delta() override{};
+    using BaseLayer::backward;
 
-    void update_weights() override{};
+    void allocate_param_delta() override {};
 
-    void update_biases() override{};
+    void update_weights() override {};
 
-    void save(std::ofstream &file) override{};
+    void update_biases() override {};
 
-    void load(std::ifstream &file) override{};
+    void save(std::ofstream &file) override {};
+
+    void load(std::ifstream &file) override {};
 
     std::unique_ptr<BaseLayer> to_host() override;
 };
@@ -383,15 +399,17 @@ class SoftmaxCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void allocate_param_delta() override{};
+    using BaseLayer::backward;
+
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
-    void save(std::ofstream &file) override{};
+    void save(std::ofstream &file) override {};
 
-    void load(std::ifstream &file) override{};
+    void load(std::ifstream &file) override {};
 
     std::unique_ptr<BaseLayer> to_host() override;
 };
\ No newline at end of file
diff --git a/include/base_layer.h b/include/base_layer.h
index c8121eb9d..8c6808359 100644
--- a/include/base_layer.h
+++ b/include/base_layer.h
@@ -98,15 +98,6 @@ class BaseLayer {
                          BaseHiddenStates &output_states,
                          BaseTempStates &temp_states);
 
-    virtual void state_backward(BaseBackwardStates &next_bwd_states,
-                                BaseDeltaStates &input_delta_states,
-                                BaseDeltaStates &output_delta_states,
-                                BaseTempStates &temp_states);
-
-    virtual void param_backward(BaseBackwardStates &next_bwd_states,
-                                BaseDeltaStates &delta_states,
-                                BaseTempStates &temp_states);
-
     virtual void backward(BaseDeltaStates &input_delta_states,
                           BaseDeltaStates &output_delta_states,
                           BaseTempStates &temp_states,
diff --git a/include/base_layer_cuda.cuh b/include/base_layer_cuda.cuh
index 2482b4455..cb2531743 100644
--- a/include/base_layer_cuda.cuh
+++ b/include/base_layer_cuda.cuh
@@ -47,8 +47,6 @@ class BaseLayerCuda : public BaseLayer {
     BaseLayerCuda &operator=(BaseLayerCuda &&) = default;
 
     using BaseLayer::forward;
-    using BaseLayer::param_backward;
-    using BaseLayer::state_backward;
 
     void allocate_param_delta() override;
 
diff --git a/include/conv2d_layer.h b/include/conv2d_layer.h
index d4a281207..a373b27af 100644
--- a/include/conv2d_layer.h
+++ b/include/conv2d_layer.h
@@ -71,15 +71,6 @@ class Conv2d : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states,
diff --git a/include/conv2d_layer_cuda.cuh b/include/conv2d_layer_cuda.cuh
index da6a51516..8c4a615f8 100644
--- a/include/conv2d_layer_cuda.cuh
+++ b/include/conv2d_layer_cuda.cuh
@@ -63,15 +63,6 @@ class Conv2dCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states,
diff --git a/include/convtranspose2d_layer.h b/include/convtranspose2d_layer.h
index d1f00cb1c..73c4997e6 100644
--- a/include/convtranspose2d_layer.h
+++ b/include/convtranspose2d_layer.h
@@ -71,14 +71,10 @@ class ConvTranspose2d : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
+    void backward(BaseDeltaStates &input_delta_states,
+                  BaseDeltaStates &output_delta_states,
+                  BaseTempStates &temp_states,
+                  bool state_udapte = true) override;
 
     using BaseLayer::storing_states_for_training;
     using BaseLayer::to_cuda;
diff --git a/include/convtranspose2d_layer_cuda.cuh b/include/convtranspose2d_layer_cuda.cuh
index 70c166f3b..c306a3dfb 100644
--- a/include/convtranspose2d_layer_cuda.cuh
+++ b/include/convtranspose2d_layer_cuda.cuh
@@ -69,14 +69,10 @@ class ConvTranspose2dCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
+    void backward(BaseDeltaStates &input_delta_states,
+                  BaseDeltaStates &output_delta_states,
+                  BaseTempStates &temp_states,
+                  bool state_udapte = true) override;
 
     std::unique_ptr<BaseLayer> to_host() override;
 
diff --git a/include/linear_layer.h b/include/linear_layer.h
index b73656e80..1eb13eb39 100644
--- a/include/linear_layer.h
+++ b/include/linear_layer.h
@@ -137,15 +137,6 @@ class Linear : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_hidden_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states, bool state_udapte) override;
diff --git a/include/linear_layer_cuda.cuh b/include/linear_layer_cuda.cuh
index 3e8042a9f..616cab0e7 100644
--- a/include/linear_layer_cuda.cuh
+++ b/include/linear_layer_cuda.cuh
@@ -94,15 +94,6 @@ class LinearCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states, bool state_udapte) override;
diff --git a/include/lstm_layer.h b/include/lstm_layer.h
index 67d198b20..e6fe17b1f 100644
--- a/include/lstm_layer.h
+++ b/include/lstm_layer.h
@@ -69,15 +69,6 @@ class LSTM : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_hidden_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states,
diff --git a/include/lstm_layer_cuda.cuh b/include/lstm_layer_cuda.cuh
index de9e379ef..df33b13a7 100644
--- a/include/lstm_layer_cuda.cuh
+++ b/include/lstm_layer_cuda.cuh
@@ -72,15 +72,6 @@ class LSTMCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states,
diff --git a/include/norm_layer.h b/include/norm_layer.h
index be98fe4a4..24ff2cae0 100644
--- a/include/norm_layer.h
+++ b/include/norm_layer.h
@@ -46,15 +46,6 @@ class LayerNorm : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states,
@@ -112,15 +103,6 @@ class BatchNorm2d : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states,
diff --git a/include/norm_layer_cuda.cuh b/include/norm_layer_cuda.cuh
index e3fdc8517..c099fe78c 100644
--- a/include/norm_layer_cuda.cuh
+++ b/include/norm_layer_cuda.cuh
@@ -45,15 +45,6 @@ class LayerNormCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states,
@@ -113,15 +104,6 @@ class BatchNorm2dCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states,
@@ -134,6 +116,7 @@ class BatchNorm2dCuda : public BaseLayerCuda {
 
    protected:
     void allocate_running_mean_var();
+    void deallocate_running_mean_var();
     void running_mean_var_to_host();
     void running_mean_var_to_device();
     void lazy_init();
diff --git a/include/pooling_layer.h b/include/pooling_layer.h
index 5bc0ff90c..276a57b6e 100644
--- a/include/pooling_layer.h
+++ b/include/pooling_layer.h
@@ -55,25 +55,16 @@ class AvgPool2d : public BaseLayer {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &ooutput_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states,
                   bool state_udapte = true) override;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
     using BaseLayer::storing_states_for_training;
     using BaseLayer::to_cuda;
diff --git a/include/pooling_layer_cuda.cuh b/include/pooling_layer_cuda.cuh
index 7b18b256c..07cb944a3 100644
--- a/include/pooling_layer_cuda.cuh
+++ b/include/pooling_layer_cuda.cuh
@@ -48,25 +48,16 @@ class AvgPool2dCuda : public BaseLayerCuda {
                  BaseHiddenStates &output_states,
                  BaseTempStates &temp_states) override;
 
-    void state_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &input_delta_states,
-                        BaseDeltaStates &output_delta_states,
-                        BaseTempStates &temp_states) override;
-
-    void param_backward(BaseBackwardStates &next_bwd_states,
-                        BaseDeltaStates &delta_states,
-                        BaseTempStates &temp_states) override;
-
     void backward(BaseDeltaStates &input_delta_states,
                   BaseDeltaStates &output_delta_states,
                   BaseTempStates &temp_states,
                   bool state_udapte = true) override;
 
-    void allocate_param_delta() override{};
+    void allocate_param_delta() override {};
 
-    void update_weights() override{};
+    void update_weights() override {};
 
-    void update_biases() override{};
+    void update_biases() override {};
 
     void preinit_layer() override;
 
diff --git a/requirements.txt b/requirements.txt
index 729af41b0..1938dd427 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,8 +2,5 @@ matplotlib==3.8.2
 ninja==1.11.1.1
 numpy==1.26.0
 pandas==2.1.4
-snakeviz==2.2.0
 tqdm==4.66.1
-fire==0.5.0
-torch==2.3.0
-torchvision==0.18.0
\ No newline at end of file
+fire==0.5.0
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 246f133b3..d14513633 100644
--- a/setup.py
+++ b/setup.py
@@ -42,7 +42,7 @@ def build_extension(self, ext):
             extdir += os.path.sep
 
         debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug
-        cfg = "Debug" if debug else "RelWithDebInfo"
+        cfg = "Debug" if debug else "Release"
 
         # CMake lets you override the generator - we need to check this.
         # Can be set with Conda-Build, for example.
diff --git a/src/activation_cuda.cu b/src/activation_cuda.cu
index 13de0ff0b..c960a5016 100644
--- a/src/activation_cuda.cu
+++ b/src/activation_cuda.cu
@@ -764,13 +764,14 @@ __global__ void mixture_sigmoid_mean_var_cuda(float const *mu_z,
         // Moments calculations (L. Alric, 2024)
         mu_a[col] = (mu_z[col] + 1) * cdf_l + (mu_z[col] - 1) * cdf_u +
                     std_z * (pdf_l - pdf_u) - mu_z[col];
-        var_a[col] = max(0.000001f,
-            (cdf_l * (var_z[col] - powf(mu_z[col], 2) - 2 * mu_z[col] - 1) +
-             cdf_u * (var_z[col] - powf(mu_z[col], 2) + 2 * mu_z[col] - 1) +
-             std_z * (pdf_u * (mu_z[col] - 1) - pdf_l * (mu_z[col] + 1)) -
-             powf(mu_a[col], 2) + 2 * mu_a[col] * mu_z[col] +
-             powf(mu_z[col], 2) - var_z[col] + 2) /
-            4.0f);
+        var_a[col] =
+            max(0.000001f,
+                (cdf_l * (var_z[col] - powf(mu_z[col], 2) - 2 * mu_z[col] - 1) +
+                 cdf_u * (var_z[col] - powf(mu_z[col], 2) + 2 * mu_z[col] - 1) +
+                 std_z * (pdf_u * (mu_z[col] - 1) - pdf_l * (mu_z[col] + 1)) -
+                 powf(mu_a[col], 2) + 2 * mu_a[col] * mu_z[col] +
+                 powf(mu_z[col], 2) - var_z[col] + 2) /
+                    4.0f);
         mu_a[col] = mu_a[col] / 2.0f + 0.5f;
         jcb[col] = (cdf_u + cdf_l - 1) / 2.0f;
     }
@@ -800,12 +801,13 @@ __global__ void mixture_tanh_mean_var_cuda(float const *mu_z,
         // Moments calculations (L. Alric, 2024)
         mu_a[col] = (mu_z[col] + 1) * cdf_l + (mu_z[col] - 1) * cdf_u +
                     std_z * (pdf_l - pdf_u) - mu_z[col];
-        var_a[col] = max(0.000001f,
+        var_a[col] = max(
+            0.000001f,
             cdf_l * (var_z[col] - powf(mu_z[col], 2) - 2 * mu_z[col] - 1) +
-            cdf_u * (var_z[col] - powf(mu_z[col], 2) + 2 * mu_z[col] - 1) +
-            std_z * (pdf_u * (mu_z[col] - 1) - pdf_l * (mu_z[col] + 1)) -
-            powf(mu_a[col], 2) + 2 * mu_a[col] * mu_z[col] +
-            powf(mu_z[col], 2) - var_z[col] + 2);
+                cdf_u * (var_z[col] - powf(mu_z[col], 2) + 2 * mu_z[col] - 1) +
+                std_z * (pdf_u * (mu_z[col] - 1) - pdf_l * (mu_z[col] + 1)) -
+                powf(mu_a[col], 2) + 2 * mu_a[col] * mu_z[col] +
+                powf(mu_z[col], 2) - var_z[col] + 2);
         jcb[col] = cdf_u + cdf_l - 1;
     }
 }
diff --git a/src/base_layer.cpp b/src/base_layer.cpp
index 732a128d6..7ef7024e3 100644
--- a/src/base_layer.cpp
+++ b/src/base_layer.cpp
@@ -49,15 +49,6 @@ void BaseLayer::forward(BaseHiddenStates &input_states,
                         BaseHiddenStates &output_states,
                         BaseTempStates &temp_states) {}
 
-void BaseLayer::state_backward(BaseBackwardStates &next_bwd_states,
-                               BaseDeltaStates &input_delta_states,
-                               BaseDeltaStates &output_hidden_states,
-                               BaseTempStates &temp_states) {}
-
-void BaseLayer::param_backward(BaseBackwardStates &next_bwd_states,
-                               BaseDeltaStates &delta_states,
-                               BaseTempStates &temp_states) {}
-
 void BaseLayer::backward(BaseDeltaStates &input_delta_states,
                          BaseDeltaStates &output_delta_states,
                          BaseTempStates &temp_states, bool state_udapte) {}
diff --git a/src/bindings/conv2d_layer_bindings.cpp b/src/bindings/conv2d_layer_bindings.cpp
index 0b2ee94f4..c5497d100 100644
--- a/src/bindings/conv2d_layer_bindings.cpp
+++ b/src/bindings/conv2d_layer_bindings.cpp
@@ -32,6 +32,5 @@ void bind_conv2d_layer(pybind11::module_& modo)
         .def_readwrite("init_method", &Conv2d::init_method)
         .def("init_weight_bias", &Conv2d::init_weight_bias)
         .def("forward", &Conv2d::forward)
-        .def("state_backward", &Conv2d::state_backward)
-        .def("param_backward", &Conv2d::param_backward);
+        .def("backward", &Conv2d::backward);
 }
\ No newline at end of file
diff --git a/src/bindings/convtranspose2d_layer_bindings.cpp b/src/bindings/convtranspose2d_layer_bindings.cpp
index b4f5ce374..24c6c13ad 100644
--- a/src/bindings/convtranspose2d_layer_bindings.cpp
+++ b/src/bindings/convtranspose2d_layer_bindings.cpp
@@ -34,6 +34,5 @@ void bind_convtranspose2d_layer(pybind11::module_& modo)
         .def_readwrite("init_method", &ConvTranspose2d::init_method)
         .def("init_weight_bias", &ConvTranspose2d::init_weight_bias)
         .def("forward", &ConvTranspose2d::forward)
-        .def("state_backward", &ConvTranspose2d::state_backward)
-        .def("param_backward", &ConvTranspose2d::param_backward);
+        .def("backward", &ConvTranspose2d::backward);
 }
\ No newline at end of file
diff --git a/src/bindings/linear_layer_bindings.cpp b/src/bindings/linear_layer_bindings.cpp
index 7690a9d20..92cda1a6f 100644
--- a/src/bindings/linear_layer_bindings.cpp
+++ b/src/bindings/linear_layer_bindings.cpp
@@ -26,8 +26,7 @@ void bind_linear_layer(pybind11::module_& modo)
         .def_readwrite("init_method", &Linear::init_method)
         .def("init_weight_bias", &Linear::init_weight_bias)
         .def("forward", &Linear::forward)
-        .def("state_backward", &Linear::state_backward)
-        .def("param_backward", &Linear::param_backward);
+        .def("state_backward", &Linear::backward);
 }
 
 // PYBIND11_MODULE(neural_net_module, m) {
diff --git a/src/bindings/lstm_layer_bindings.cpp b/src/bindings/lstm_layer_bindings.cpp
index dc2ead911..d17329925 100644
--- a/src/bindings/lstm_layer_bindings.cpp
+++ b/src/bindings/lstm_layer_bindings.cpp
@@ -34,6 +34,5 @@ void bind_lstm_layer(pybind11::module_& modo)
         .def_readwrite("init_method", &LSTM::init_method)
         .def("init_weight_bias", &LSTM::init_weight_bias)
         .def("forward", &LSTM::forward)
-        .def("state_backward", &LSTM::state_backward)
-        .def("param_backward", &LSTM::param_backward);
+        .def("backward", &LSTM::backward);
 }
\ No newline at end of file
diff --git a/src/bindings/norm_layer_bindings.cpp b/src/bindings/norm_layer_bindings.cpp
index 354459a80..3702c0bc8 100644
--- a/src/bindings/norm_layer_bindings.cpp
+++ b/src/bindings/norm_layer_bindings.cpp
@@ -27,8 +27,7 @@ void bind_layernorm_layer(pybind11::module_& modo)
         .def("get_layer_info", &LayerNorm::get_layer_info)
         .def("get_layer_name", &LayerNorm::get_layer_name)
         .def("forward", &LayerNorm::forward)
-        .def("state_backward", &LayerNorm::state_backward)
-        .def("param_backward", &LayerNorm::param_backward)
+        .def("backward", &LayerNorm::backward)
         .def("init_weight_bias", &LayerNorm::init_weight_bias);
     ;
 }
@@ -45,8 +44,7 @@ void bind_batchnorm_layer(pybind11::module_& modo)
         .def("get_layer_info", &BatchNorm2d::get_layer_info)
         .def("get_layer_name", &BatchNorm2d::get_layer_name)
         .def("forward", &BatchNorm2d::forward)
-        .def("state_backward", &BatchNorm2d::state_backward)
-        .def("param_backward", &BatchNorm2d::param_backward)
+        .def("backward", &BatchNorm2d::backward)
         .def("init_weight_bias", &BatchNorm2d::init_weight_bias);
     ;
 }
\ No newline at end of file
diff --git a/src/bindings/pooling_layer_bindings.cpp b/src/bindings/pooling_layer_bindings.cpp
index 4737cb9ec..c3f72356d 100644
--- a/src/bindings/pooling_layer_bindings.cpp
+++ b/src/bindings/pooling_layer_bindings.cpp
@@ -27,7 +27,6 @@ void bind_avgpool2d_layer(pybind11::module_& modo)
         .def("get_layer_info", &AvgPool2d::get_layer_info)
         .def("get_layer_name", &AvgPool2d::get_layer_name)
         .def("forward", &AvgPool2d::forward)
-        .def("state_backward", &AvgPool2d::state_backward)
-        .def("param_backward", &AvgPool2d::param_backward);
+        .def("state_backward", &AvgPool2d::backward);
     ;
 }
\ No newline at end of file
diff --git a/src/conv2d_layer.cpp b/src/conv2d_layer.cpp
index 036bb97ba..2dc6229cb 100644
--- a/src/conv2d_layer.cpp
+++ b/src/conv2d_layer.cpp
@@ -179,85 +179,6 @@ void Conv2d::forward(BaseHiddenStates &input_states,
     }
 }
 
-void Conv2d::state_backward(BaseBackwardStates &next_bwd_states,
-                            BaseDeltaStates &input_delta_states,
-                            BaseDeltaStates &output_delta_states,
-                            BaseTempStates &temp_states)
-/*
- */
-{
-    // Initialization
-    int batch_size = input_delta_states.block_size;
-
-    // Launch kernel
-    int wihi = this->in_width * this->in_height;
-    int woho = this->out_width * this->out_height;
-    int row_zw_fo = this->row_zw * this->out_channels;
-    int pad_idx = woho * this->out_channels * batch_size + 1;
-
-    permute_jacobian(next_bwd_states.jcb, wihi, this->in_channels, batch_size,
-                     temp_states.tmp_1);
-
-    if (this->num_threads > 1) {
-        conv2d_bwd_delta_z_mp(
-            this->mu_w, temp_states.tmp_1, input_delta_states.delta_mu,
-            input_delta_states.delta_var, this->idx_cov_zwa_1,
-            this->idx_var_z_ud, woho, this->out_channels, wihi,
-            this->in_channels, this->kernel_size, this->row_zw, row_zw_fo,
-            batch_size, pad_idx, this->num_threads,
-            output_delta_states.delta_mu, output_delta_states.delta_var);
-    } else {
-        conv2d_bwd_delta_z(
-            this->mu_w, temp_states.tmp_1, input_delta_states.delta_mu,
-            input_delta_states.delta_var, this->idx_cov_zwa_1,
-            this->idx_var_z_ud, woho, this->out_channels, wihi,
-            this->in_channels, this->kernel_size, this->row_zw, row_zw_fo,
-            batch_size, pad_idx, 0, wihi * batch_size * this->in_channels,
-            output_delta_states.delta_mu, output_delta_states.delta_var);
-    }
-}
-
-void Conv2d::param_backward(BaseBackwardStates &next_bwd_states,
-                            BaseDeltaStates &delta_states,
-                            BaseTempStates &temp_states)
-/*
- */
-{
-    // Initalization
-    int batch_size = delta_states.block_size;
-
-    int woho = this->out_width * this->out_height;
-    int wihi = this->in_width * this->in_height;
-    int woho_batch = woho * batch_size;
-    int wohofo = woho * this->out_channels;
-    int pad_idx = wihi * this->in_channels * batch_size + 1;
-
-    permute_delta(delta_states.delta_mu, delta_states.delta_var, woho, wohofo,
-                  batch_size, temp_states.tmp_1, temp_states.tmp_2);
-
-    if (this->num_threads > 1) {
-        conv2d_bwd_delta_w_mp(
-            this->var_w, next_bwd_states.mu_a, temp_states.tmp_1,
-            temp_states.tmp_2, this->idx_mwa_2, batch_size, this->out_channels,
-            woho, wihi, this->in_channels, this->kernel_size, pad_idx,
-            this->num_threads, this->delta_mu_w, this->delta_var_w);
-    } else {
-        int end_chunk = this->kernel_size * this->kernel_size *
-                        this->in_channels * this->out_channels;
-        conv2d_bwd_delta_w(this->var_w, next_bwd_states.mu_a, temp_states.tmp_1,
-                           temp_states.tmp_2, this->idx_mwa_2, batch_size,
-                           this->out_channels, woho, wihi, this->in_channels,
-                           this->kernel_size, pad_idx, 0, end_chunk,
-                           this->delta_mu_w, this->delta_var_w);
-    }
-
-    if (this->bias) {
-        conv2d_bwd_delta_b(this->var_b, temp_states.tmp_1, temp_states.tmp_2,
-                           woho_batch, this->out_channels, this->delta_mu_b,
-                           this->delta_var_b);
-    }
-}
-
 void Conv2d::backward(BaseDeltaStates &input_delta_states,
                       BaseDeltaStates &output_delta_states,
                       BaseTempStates &temp_states, bool state_udapte)
diff --git a/src/conv2d_layer_cuda.cu b/src/conv2d_layer_cuda.cu
index 48c0f1942..7d6783f3d 100644
--- a/src/conv2d_layer_cuda.cu
+++ b/src/conv2d_layer_cuda.cu
@@ -517,8 +517,8 @@ __global__ void conv2d_bwd_delta_z_cuda_tiled(
         }
 
         __syncthreads();
-#pragma unroll
         float tmp_mu_w = 0;
+#pragma unroll
         for (int i = 0; i < TILE_SIZE; i++) {
             int tmp_zw_idx = s_zw_idx[tx * PADDED_TILE_SIZE + i];
 
@@ -858,111 +858,6 @@ void Conv2dCuda::forward(BaseHiddenStates &input_states,
     }
 }
 
-void Conv2dCuda::state_backward(BaseBackwardStates &next_bwd_states,
-                                BaseDeltaStates &input_delta_states,
-                                BaseDeltaStates &output_delta_states,
-                                BaseTempStates &temp_states)
-/*
- */
-{
-    // New poitner will point to the same memory location when casting
-    BackwardStateCuda *cu_next_bwd_states =
-        dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
-    DeltaStateCuda *cu_input_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&input_delta_states);
-    DeltaStateCuda *cu_output_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&output_delta_states);
-    TempStateCuda *cu_temp_states = dynamic_cast<TempStateCuda *>(&temp_states);
-
-    // Initialization
-    int batch_size = input_delta_states.block_size;
-    int threads = this->num_cuda_threads;
-
-    // Launch kernel
-    int wihi = this->in_width * this->in_height;
-    int woho = this->out_width * this->out_height;
-    int row_zw_fo = this->row_zw * this->out_channels;
-    int pad_idx = woho * this->out_channels * batch_size + 1;
-
-    unsigned int grid_row_p = (batch_size + threads - 1) / threads;
-    unsigned int grid_col_p =
-        (wihi * this->in_channels + threads - 1) / threads;
-    dim3 dim_grid_p(grid_col_p, grid_row_p);
-
-    unsigned int grid_row = (this->in_channels + threads - 1) / threads;
-    unsigned int grid_col = (wihi * batch_size + threads - 1) / threads;
-    dim3 dim_grid(grid_col, grid_row);
-    dim3 dim_block(threads, threads);
-
-    permmute_jacobian_cuda<<<dim_grid_p, dim_block>>>(
-        cu_next_bwd_states->d_jcb, wihi, this->in_channels, batch_size,
-        cu_temp_states->d_tmp_1);
-
-    conv2d_bwd_delta_z_cuda<<<dim_grid, dim_block>>>(
-        this->d_mu_w, cu_temp_states->d_tmp_1,
-        cu_input_delta_states->d_delta_mu, cu_input_delta_states->d_delta_var,
-        this->d_idx_cov_zwa_1, this->d_idx_var_z_ud, woho, this->out_channels,
-        wihi, this->in_channels, this->kernel_size, this->row_zw, row_zw_fo,
-        batch_size, pad_idx, cu_output_delta_states->d_delta_mu,
-        cu_output_delta_states->d_delta_var);
-}
-
-void Conv2dCuda::param_backward(BaseBackwardStates &next_bwd_states,
-                                BaseDeltaStates &delta_states,
-                                BaseTempStates &temp_states)
-/*
- */
-{
-    // New poitner will point to the same memory location when casting
-    BackwardStateCuda *cu_next_bwd_states =
-        dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
-    DeltaStateCuda *cu_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&delta_states);
-    TempStateCuda *cu_temp_states = dynamic_cast<TempStateCuda *>(&temp_states);
-
-    // Initalization
-    int batch_size = delta_states.block_size;
-    int threads = this->num_cuda_threads;
-
-    // Lauch kernel
-    int woho = this->out_width * this->out_height;
-    int wihi = this->in_width * this->in_height;
-    int woho_batch = woho * batch_size;
-    int wohofo = woho * this->out_channels;
-    int pad_idx = wihi * this->in_channels * batch_size + 1;
-    int ki2_fi = this->kernel_size * this->kernel_size * this->in_channels;
-
-    unsigned int grid_row = (batch_size + threads - 1) / threads;
-    unsigned int grid_col = (wohofo + threads - 1) / threads;
-    unsigned int grid_row_w = (ki2_fi + threads - 1) / threads;
-    unsigned int grid_col_w = (this->out_channels + threads - 1) / threads;
-
-    dim3 dim_grid(grid_col, grid_row);
-    dim3 dim_grid_w(grid_col_w, grid_row_w);
-    dim3 dim_block(threads, threads);
-
-    permute_delta_cuda<<<dim_grid, dim_block>>>(
-        cu_delta_states->d_delta_mu, cu_delta_states->d_delta_var, woho, wohofo,
-        batch_size, cu_temp_states->d_tmp_1, cu_temp_states->d_tmp_2);
-
-    conv2d_bwd_delta_w_cuda<<<dim_grid_w, dim_block>>>(
-        this->d_var_w, cu_next_bwd_states->d_mu_a, cu_temp_states->d_tmp_1,
-        cu_temp_states->d_tmp_2, this->d_idx_mwa_2, batch_size,
-        this->out_channels, woho, wihi, this->in_channels, this->kernel_size,
-        pad_idx, this->d_delta_mu_w, this->d_delta_var_w);
-
-    if (this->bias) {
-        unsigned int grid_col_bias =
-            (this->out_channels + threads - 1) / threads;
-        // dim3 dim_grid_bias(grid_col_bias, 1);
-
-        conv2d_bwd_delta_b_cuda<<<grid_col_bias, threads>>>(
-            this->d_var_b, cu_temp_states->d_tmp_1, cu_temp_states->d_tmp_2,
-            woho_batch, this->out_channels, this->d_delta_mu_b,
-            this->d_delta_var_b);
-    }
-}
-
 void Conv2dCuda::backward(BaseDeltaStates &input_delta_states,
                           BaseDeltaStates &output_delta_states,
                           BaseTempStates &temp_states, bool state_udapte)
@@ -989,7 +884,7 @@ void Conv2dCuda::backward(BaseDeltaStates &input_delta_states,
     int pad_param_idx = this->num_weights + 1;
 
     if (state_udapte) {
-        int pad_idx = woho * this->out_channels * batch_size + 1;
+        // int pad_idx = woho * this->out_channels * batch_size + 1;
 
         unsigned int grid_row_p = (batch_size + threads - 1) / threads;
         unsigned int grid_col_p =
diff --git a/src/convtranspose2d_layer.cpp b/src/convtranspose2d_layer.cpp
index 89747b485..e9db54048 100644
--- a/src/convtranspose2d_layer.cpp
+++ b/src/convtranspose2d_layer.cpp
@@ -615,10 +615,9 @@ void ConvTranspose2d::forward(BaseHiddenStates &input_states,
     }
 }
 
-void ConvTranspose2d::state_backward(BaseBackwardStates &next_bwd_states,
-                                     BaseDeltaStates &input_delta_states,
-                                     BaseDeltaStates &output_delta_states,
-                                     BaseTempStates &temp_states)
+void ConvTranspose2d::backward(BaseDeltaStates &input_delta_states,
+                               BaseDeltaStates &output_delta_states,
+                               BaseTempStates &temp_states, bool state_udapte)
 /*
  */
 {
@@ -630,7 +629,7 @@ void ConvTranspose2d::state_backward(BaseBackwardStates &next_bwd_states,
 
     if (this->num_threads > 1) {
         convtranspose2d_bwd_delta_z_mp(
-            this->mu_w, next_bwd_states.jcb, input_delta_states.delta_mu,
+            this->mu_w, this->bwd_states->jcb, input_delta_states.delta_mu,
             input_delta_states.delta_var, this->idx_cov_z_wa_1,
             this->idx_var_z_ud, woho, this->out_channels, wihi,
             this->in_channels, this->kernel_size, this->row_zw, batch_size,
@@ -638,55 +637,121 @@ void ConvTranspose2d::state_backward(BaseBackwardStates &next_bwd_states,
             output_delta_states.delta_var);
     } else {
         convtranspose2d_bwd_delta_z(
-            this->mu_w, next_bwd_states.jcb, input_delta_states.delta_mu,
+            this->mu_w, this->bwd_states->jcb, input_delta_states.delta_mu,
             input_delta_states.delta_var, this->idx_cov_z_wa_1,
             this->idx_var_z_ud, woho, this->out_channels, wihi,
             this->in_channels, this->kernel_size, this->row_zw, 0, batch_size,
             output_delta_states.delta_mu, output_delta_states.delta_var);
     }
-}
 
-void ConvTranspose2d::param_backward(BaseBackwardStates &next_bwd_states,
-                                     BaseDeltaStates &delta_states,
-                                     BaseTempStates &temp_states)
-/*
- */
-{
-    int batch_size = delta_states.block_size;
+    // Parameters
 
     int ki2 = this->kernel_size * this->kernel_size;
-    int wihi = this->in_height * this->in_width;
-    int woho = this->out_width * this->out_height;
 
     if (this->num_threads > 1) {
         convtranspose2d_bwd_delta_w_mp(
-            this->var_w, next_bwd_states.mu_a, delta_states.delta_mu,
-            delta_states.delta_var, this->idx_cov_wz_2, this->idx_var_wz_ud,
-            woho, this->out_channels, wihi, this->in_channels,
-            this->kernel_size, batch_size, this->num_threads, this->delta_mu_w,
-            this->delta_var_w);
+            this->var_w, this->bwd_states->mu_a, input_delta_states.delta_mu,
+            input_delta_states.delta_var, this->idx_cov_wz_2,
+            this->idx_var_wz_ud, woho, this->out_channels, wihi,
+            this->in_channels, this->kernel_size, batch_size, this->num_threads,
+            this->delta_mu_w, this->delta_var_w);
         if (this->bias) {
             convtranspose2d_bwd_delta_b_mp(
-                this->var_b, delta_states.delta_mu, delta_states.delta_var,
-                woho, this->out_channels, batch_size, this->num_threads,
-                this->delta_mu_b, this->delta_var_b);
+                this->var_b, input_delta_states.delta_mu,
+                input_delta_states.delta_var, woho, this->out_channels,
+                batch_size, this->num_threads, this->delta_mu_b,
+                this->delta_var_b);
         }
     } else {
         convtranspose2d_bwd_delta_w(
-            this->var_w, next_bwd_states.mu_a, delta_states.delta_mu,
-            delta_states.delta_var, this->idx_cov_wz_2, this->idx_var_wz_ud,
-            woho, this->out_channels, wihi, this->in_channels,
-            this->kernel_size, batch_size, 0, ki2 * this->out_channels,
-            this->delta_mu_w, this->delta_var_w);
+            this->var_w, this->bwd_states->mu_a, input_delta_states.delta_mu,
+            input_delta_states.delta_var, this->idx_cov_wz_2,
+            this->idx_var_wz_ud, woho, this->out_channels, wihi,
+            this->in_channels, this->kernel_size, batch_size, 0,
+            ki2 * this->out_channels, this->delta_mu_w, this->delta_var_w);
         if (this->bias) {
             convtranspose2d_bwd_delta_b(
-                this->var_b, delta_states.delta_mu, delta_states.delta_var,
-                woho, this->out_channels, batch_size, 0, this->out_channels,
-                this->delta_mu_b, this->delta_var_b);
+                this->var_b, input_delta_states.delta_mu,
+                input_delta_states.delta_var, woho, this->out_channels,
+                batch_size, 0, this->out_channels, this->delta_mu_b,
+                this->delta_var_b);
         }
     }
 }
 
+// void ConvTranspose2d::state_backward(BaseBackwardStates &next_bwd_states,
+//                                      BaseDeltaStates &input_delta_states,
+//                                      BaseDeltaStates &output_delta_states,
+//                                      BaseTempStates &temp_states)
+// /*
+//  */
+// {
+//     // Initialization
+//     int batch_size = input_delta_states.block_size;
+
+//     int wihi = this->in_height * this->in_width;
+//     int woho = this->out_width * this->out_height;
+
+//     if (this->num_threads > 1) {
+//         convtranspose2d_bwd_delta_z_mp(
+//             this->mu_w, next_bwd_states.jcb, input_delta_states.delta_mu,
+//             input_delta_states.delta_var, this->idx_cov_z_wa_1,
+//             this->idx_var_z_ud, woho, this->out_channels, wihi,
+//             this->in_channels, this->kernel_size, this->row_zw, batch_size,
+//             this->num_threads, output_delta_states.delta_mu,
+//             output_delta_states.delta_var);
+//     } else {
+//         convtranspose2d_bwd_delta_z(
+//             this->mu_w, next_bwd_states.jcb, input_delta_states.delta_mu,
+//             input_delta_states.delta_var, this->idx_cov_z_wa_1,
+//             this->idx_var_z_ud, woho, this->out_channels, wihi,
+//             this->in_channels, this->kernel_size, this->row_zw, 0,
+//             batch_size, output_delta_states.delta_mu,
+//             output_delta_states.delta_var);
+//     }
+// }
+
+// void ConvTranspose2d::param_backward(BaseBackwardStates &next_bwd_states,
+//                                      BaseDeltaStates &delta_states,
+//                                      BaseTempStates &temp_states)
+// /*
+//  */
+// {
+//     int batch_size = delta_states.block_size;
+
+//     int ki2 = this->kernel_size * this->kernel_size;
+//     int wihi = this->in_height * this->in_width;
+//     int woho = this->out_width * this->out_height;
+
+//     if (this->num_threads > 1) {
+//         convtranspose2d_bwd_delta_w_mp(
+//             this->var_w, next_bwd_states.mu_a, delta_states.delta_mu,
+//             delta_states.delta_var, this->idx_cov_wz_2, this->idx_var_wz_ud,
+//             woho, this->out_channels, wihi, this->in_channels,
+//             this->kernel_size, batch_size, this->num_threads,
+//             this->delta_mu_w, this->delta_var_w);
+//         if (this->bias) {
+//             convtranspose2d_bwd_delta_b_mp(
+//                 this->var_b, delta_states.delta_mu, delta_states.delta_var,
+//                 woho, this->out_channels, batch_size, this->num_threads,
+//                 this->delta_mu_b, this->delta_var_b);
+//         }
+//     } else {
+//         convtranspose2d_bwd_delta_w(
+//             this->var_w, next_bwd_states.mu_a, delta_states.delta_mu,
+//             delta_states.delta_var, this->idx_cov_wz_2, this->idx_var_wz_ud,
+//             woho, this->out_channels, wihi, this->in_channels,
+//             this->kernel_size, batch_size, 0, ki2 * this->out_channels,
+//             this->delta_mu_w, this->delta_var_w);
+//         if (this->bias) {
+//             convtranspose2d_bwd_delta_b(
+//                 this->var_b, delta_states.delta_mu, delta_states.delta_var,
+//                 woho, this->out_channels, batch_size, 0, this->out_channels,
+//                 this->delta_mu_b, this->delta_var_b);
+//         }
+//     }
+// }
+
 #ifdef USE_CUDA
 std::unique_ptr<BaseLayer> ConvTranspose2d::to_cuda() {
     this->device = "cuda";
diff --git a/src/convtranspose2d_layer_cuda.cu b/src/convtranspose2d_layer_cuda.cu
index dc8b01b23..5803ed653 100644
--- a/src/convtranspose2d_layer_cuda.cu
+++ b/src/convtranspose2d_layer_cuda.cu
@@ -402,16 +402,16 @@ void ConvTranspose2dCuda::forward(BaseHiddenStates &input_states,
     }
 }
 
-void ConvTranspose2dCuda::state_backward(BaseBackwardStates &next_bwd_states,
-                                         BaseDeltaStates &input_delta_states,
-                                         BaseDeltaStates &output_delta_states,
-                                         BaseTempStates &temp_states)
+void ConvTranspose2dCuda::backward(BaseDeltaStates &input_delta_states,
+                                   BaseDeltaStates &output_delta_states,
+                                   BaseTempStates &temp_states,
+                                   bool state_udapte)
 /*
  */
 {
     // New poitner will point to the same memory location when casting
     BackwardStateCuda *cu_next_bwd_states =
-        dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
+        dynamic_cast<BackwardStateCuda *>(this->bwd_states.get());
     DeltaStateCuda *cu_input_delta_states =
         dynamic_cast<DeltaStateCuda *>(&input_delta_states);
     DeltaStateCuda *cu_output_delta_states =
@@ -436,51 +436,117 @@ void ConvTranspose2dCuda::state_backward(BaseBackwardStates &next_bwd_states,
         wihi, this->in_channels, this->kernel_size, this->row_zw, batch_size,
         cu_output_delta_states->d_delta_mu,
         cu_output_delta_states->d_delta_var);
-}
-
-void ConvTranspose2dCuda::param_backward(BaseBackwardStates &next_bwd_states,
-                                         BaseDeltaStates &delta_states,
-                                         BaseTempStates &temp_states)
-/**/
-{
-    // New poitner will point to the same memory location when casting
-    BackwardStateCuda *cu_next_bwd_states =
-        dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
-    DeltaStateCuda *cu_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&delta_states);
 
-    // Initalization
-    int batch_size = delta_states.block_size;
-    int threads = this->num_cuda_threads;
+    // Parameter
 
     // Launch kernel
     int ki2 = this->kernel_size * this->kernel_size;
-    int wihi = this->in_height * this->in_width;
-    int woho = this->out_width * this->out_height;
     unsigned int grid_row_w = (this->in_channels + threads - 1) / threads;
     unsigned int grid_col_w =
         (ki2 * this->out_channels + threads - 1) / threads;
 
     dim3 dim_grid_w(grid_col_w, grid_row_w);
-    dim3 dim_block(threads, threads);
 
     convtranspose2d_bwd_delta_w_cuda<<<dim_grid_w, dim_block>>>(
-        this->d_var_w, cu_next_bwd_states->d_mu_a, cu_delta_states->d_delta_mu,
-        cu_delta_states->d_delta_var, this->d_idx_cov_wz_2,
-        this->d_idx_var_wz_ud, woho, this->out_channels, wihi,
-        this->in_channels, this->kernel_size, batch_size, this->d_delta_mu_w,
-        this->d_delta_var_w);
+        this->d_var_w, cu_next_bwd_states->d_mu_a,
+        cu_input_delta_states->d_delta_mu, cu_input_delta_states->d_delta_var,
+        this->d_idx_cov_wz_2, this->d_idx_var_wz_ud, woho, this->out_channels,
+        wihi, this->in_channels, this->kernel_size, batch_size,
+        this->d_delta_mu_w, this->d_delta_var_w);
 
     if (this->bias) {
         unsigned int grid_size = (this->out_channels + threads - 1) / threads;
 
         convtranspose2d_bwd_delta_b_cuda<<<grid_size, threads>>>(
-            this->d_var_b, cu_delta_states->d_delta_mu,
-            cu_delta_states->d_delta_var, woho, this->out_channels, batch_size,
-            this->d_delta_mu_b, this->d_delta_var_b);
+            this->d_var_b, cu_input_delta_states->d_delta_mu,
+            cu_input_delta_states->d_delta_var, woho, this->out_channels,
+            batch_size, this->d_delta_mu_b, this->d_delta_var_b);
     }
 }
 
+// void ConvTranspose2dCuda::state_backward(BaseBackwardStates &next_bwd_states,
+//                                          BaseDeltaStates &input_delta_states,
+//                                          BaseDeltaStates
+//                                          &output_delta_states, BaseTempStates
+//                                          &temp_states)
+// /*
+//  */
+// {
+//     // New poitner will point to the same memory location when casting
+//     BackwardStateCuda *cu_next_bwd_states =
+//         dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
+//     DeltaStateCuda *cu_input_delta_states =
+//         dynamic_cast<DeltaStateCuda *>(&input_delta_states);
+//     DeltaStateCuda *cu_output_delta_states =
+//         dynamic_cast<DeltaStateCuda *>(&output_delta_states);
+
+//     int batch_size = input_delta_states.block_size;
+//     int threads = this->num_cuda_threads;
+
+//     // Lauch kernel
+//     int wihi = this->in_height * this->in_width;
+//     int woho = this->out_width * this->out_height;
+//     unsigned int grid_row = (batch_size + threads - 1) / threads;
+//     unsigned int grid_col = (wihi * this->in_channels + threads - 1) /
+//     threads;
+
+//     dim3 dim_grid(grid_col, grid_row);
+//     dim3 dim_block(threads, threads);
+
+//     convtranspose2d_bwd_delta_z_cuda<<<dim_grid, dim_block>>>(
+//         this->d_mu_w, cu_next_bwd_states->d_jcb,
+//         cu_input_delta_states->d_delta_mu,
+//         cu_input_delta_states->d_delta_var, this->d_idx_cov_z_wa_1,
+//         this->d_idx_var_z_ud, woho, this->out_channels, wihi,
+//         this->in_channels, this->kernel_size, this->row_zw, batch_size,
+//         cu_output_delta_states->d_delta_mu,
+//         cu_output_delta_states->d_delta_var);
+// }
+
+// void ConvTranspose2dCuda::param_backward(BaseBackwardStates &next_bwd_states,
+//                                          BaseDeltaStates &delta_states,
+//                                          BaseTempStates &temp_states)
+// /**/
+// {
+//     // New poitner will point to the same memory location when casting
+//     BackwardStateCuda *cu_next_bwd_states =
+//         dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
+//     DeltaStateCuda *cu_delta_states =
+//         dynamic_cast<DeltaStateCuda *>(&delta_states);
+
+//     // Initalization
+//     int batch_size = delta_states.block_size;
+//     int threads = this->num_cuda_threads;
+
+//     // Launch kernel
+//     int ki2 = this->kernel_size * this->kernel_size;
+//     int wihi = this->in_height * this->in_width;
+//     int woho = this->out_width * this->out_height;
+//     unsigned int grid_row_w = (this->in_channels + threads - 1) / threads;
+//     unsigned int grid_col_w =
+//         (ki2 * this->out_channels + threads - 1) / threads;
+
+//     dim3 dim_grid_w(grid_col_w, grid_row_w);
+//     dim3 dim_block(threads, threads);
+
+//     convtranspose2d_bwd_delta_w_cuda<<<dim_grid_w, dim_block>>>(
+//         this->d_var_w, cu_next_bwd_states->d_mu_a,
+//         cu_delta_states->d_delta_mu, cu_delta_states->d_delta_var,
+//         this->d_idx_cov_wz_2, this->d_idx_var_wz_ud, woho,
+//         this->out_channels, wihi, this->in_channels, this->kernel_size,
+//         batch_size, this->d_delta_mu_w, this->d_delta_var_w);
+
+//     if (this->bias) {
+//         unsigned int grid_size = (this->out_channels + threads - 1) /
+//         threads;
+
+//         convtranspose2d_bwd_delta_b_cuda<<<grid_size, threads>>>(
+//             this->d_var_b, cu_delta_states->d_delta_mu,
+//             cu_delta_states->d_delta_var, woho, this->out_channels,
+//             batch_size, this->d_delta_mu_b, this->d_delta_var_b);
+//     }
+// }
+
 std::unique_ptr<BaseLayer> ConvTranspose2dCuda::to_host()
 /* Transfer to cpu version
  */
diff --git a/src/data_struct_cuda.cu b/src/data_struct_cuda.cu
index 9b9dc1bda..70103a648 100644
--- a/src/data_struct_cuda.cu
+++ b/src/data_struct_cuda.cu
@@ -465,6 +465,9 @@ void BackwardStateCuda::allocate_memory()
 /*
  */
 {
+    if (this->d_mu_a != nullptr || this->d_jcb != nullptr) {
+        this->deallocate_memory();
+    }
     this->mu_a.resize(this->size, 0);
     this->jcb.resize(this->size, 0);
     cudaMalloc(&this->d_mu_a, this->size * sizeof(float));
diff --git a/src/debugger.cpp b/src/debugger.cpp
index 8b94c27ae..eda0f4b89 100644
--- a/src/debugger.cpp
+++ b/src/debugger.cpp
@@ -149,8 +149,8 @@ void ModelDebugger::debug_forward(const std::vector<float> &mu_x,
             if (this->test_output_z_buffer->mu_a[j] !=
                 this->ref_output_z_buffer->mu_a[j]) {
                 auto layer_name = test_current_layer->get_layer_name();
-                std::cout << "Layer name: " << layer_name << " "
-                          << "Layer no " << i << "\n"
+                std::cout << "Layer name: " << layer_name << " " << "Layer no "
+                          << i << "\n"
                           << std::endl;
                 // std::vector<float> test_mu_ra, test_var_ra, ref_mu_ra,
                 //     ref_var_ra;
@@ -242,25 +242,25 @@ void ModelDebugger::debug_backward(std::vector<float> &y_batch,
         auto *test_current_layer = test_model.layers[i].get();
         auto *ref_current_layer = ref_model.layers[i].get();
 
-        // Backward pass for parameters and hidden states
-        if (test_model.param_update) {
-            test_current_layer->param_backward(*test_current_layer->bwd_states,
-                                               *test_input_delta_z_buffer,
-                                               *test_temp_states);
+        // // Backward pass for parameters and hidden states
+        // if (test_model.param_update) {
+        //     test_current_layer->param_backward(*test_current_layer->bwd_states,
+        //                                        *test_input_delta_z_buffer,
+        //                                        *test_temp_states);
 
-            ref_current_layer->param_backward(*ref_current_layer->bwd_states,
-                                              *ref_input_delta_z_buffer,
-                                              *ref_temp_states);
-        }
+        //     ref_current_layer->param_backward(*ref_current_layer->bwd_states,
+        //                                       *ref_input_delta_z_buffer,
+        //                                       *ref_temp_states);
+        // }
 
-        // Backward pass for hidden states
-        test_current_layer->state_backward(
-            *test_current_layer->bwd_states, *test_input_delta_z_buffer,
-            *test_output_delta_z_buffer, *test_temp_states);
+        // // Backward pass for hidden states
+        // test_current_layer->state_backward(
+        //     *test_current_layer->bwd_states, *test_input_delta_z_buffer,
+        //     *test_output_delta_z_buffer, *test_temp_states);
 
-        ref_current_layer->state_backward(
-            *ref_current_layer->bwd_states, *ref_input_delta_z_buffer,
-            *ref_output_delta_z_buffer, *ref_temp_states);
+        // ref_current_layer->state_backward(
+        //     *ref_current_layer->bwd_states, *ref_input_delta_z_buffer,
+        //     *ref_output_delta_z_buffer, *ref_temp_states);
 
         // Copy to host for gpu model
 #ifdef USE_CUDA
@@ -296,9 +296,8 @@ void ModelDebugger::debug_backward(std::vector<float> &y_batch,
                 if (this->test_output_delta_z_buffer->delta_var[j] !=
                     this->ref_output_delta_z_buffer->delta_var[j]) {
                     std::cout << "Layer name: " << layer_name << " "
-                              << "Hidden states "
-                              << " "
-                              << "Layer no " << i << "\n"
+                              << "Hidden states " << " " << "Layer no " << i
+                              << "\n"
                               << std::endl;
 
                     int check = 0;
@@ -309,9 +308,7 @@ void ModelDebugger::debug_backward(std::vector<float> &y_batch,
                 if (test_current_layer->delta_var_w[k] !=
                     ref_current_layer->delta_var_w[k]) {
                     std::cout << "Layer name: " << layer_name << " "
-                              << "Weight "
-                              << " "
-                              << "Layer no " << i << "\n"
+                              << "Weight " << " " << "Layer no " << i << "\n"
                               << std::endl;
 
                     int check = 0;
@@ -345,31 +342,31 @@ void ModelDebugger::debug_backward(std::vector<float> &y_batch,
         }
     }
 
-    // Parameter update for input layer
-    if (test_model.param_update) {
-        test_model.layers[0]->param_backward(*test_model.layers[0]->bwd_states,
-                                             *test_input_delta_z_buffer,
-                                             *test_temp_states);
-    }
-
-    // State update for input layer
-    if (test_model.input_state_update) {
-        test_model.layers[0]->state_backward(
-            *test_model.layers[0]->bwd_states, *test_input_delta_z_buffer,
-            *test_output_delta_z_buffer, *test_temp_states);
-    }
-
-    if (ref_model.param_update) {
-        ref_model.layers[0]->param_backward(*ref_model.layers[0]->bwd_states,
-                                            *ref_input_delta_z_buffer,
-                                            *ref_temp_states);
-    }
-
-    if (ref_model.input_state_update) {
-        ref_model.layers[0]->state_backward(
-            *ref_model.layers[0]->bwd_states, *ref_input_delta_z_buffer,
-            *ref_output_delta_z_buffer, *ref_temp_states);
-    }
+    // // Parameter update for input layer
+    // if (test_model.param_update) {
+    //     test_model.layers[0]->param_backward(*test_model.layers[0]->bwd_states,
+    //                                          *test_input_delta_z_buffer,
+    //                                          *test_temp_states);
+    // }
+
+    // // State update for input layer
+    // if (test_model.input_state_update) {
+    //     test_model.layers[0]->state_backward(
+    //         *test_model.layers[0]->bwd_states, *test_input_delta_z_buffer,
+    //         *test_output_delta_z_buffer, *test_temp_states);
+    // }
+
+    // if (ref_model.param_update) {
+    //     ref_model.layers[0]->param_backward(*ref_model.layers[0]->bwd_states,
+    //                                         *ref_input_delta_z_buffer,
+    //                                         *ref_temp_states);
+    // }
+
+    // if (ref_model.input_state_update) {
+    //     ref_model.layers[0]->state_backward(
+    //         *ref_model.layers[0]->bwd_states, *ref_input_delta_z_buffer,
+    //         *ref_output_delta_z_buffer, *ref_temp_states);
+    // }
 }
 
 #ifdef USE_CUDA
@@ -508,17 +505,17 @@ void CrossValidator::validate_backward(std::vector<float> &y_batch,
     for (int i = num_layers - 1; i > 0; i--) {
         auto *test_current_layer = test_model.layers[i].get();
 
-        // Backward pass for parameters and hidden states
-        if (test_model.param_update) {
-            test_current_layer->param_backward(*test_current_layer->bwd_states,
-                                               *test_input_delta_z_buffer,
-                                               *test_temp_states);
-        }
+        // // Backward pass for parameters and hidden states
+        // if (test_model.param_update) {
+        //     test_current_layer->param_backward(*test_current_layer->bwd_states,
+        //                                        *test_input_delta_z_buffer,
+        //                                        *test_temp_states);
+        // }
 
-        // Backward pass for hidden states
-        test_current_layer->state_backward(
-            *test_current_layer->bwd_states, *test_input_delta_z_buffer,
-            *test_output_delta_z_buffer, *test_temp_states);
+        // // Backward pass for hidden states
+        // test_current_layer->state_backward(
+        //     *test_current_layer->bwd_states, *test_input_delta_z_buffer,
+        //     *test_output_delta_z_buffer, *test_temp_states);
 
         // Copy to host for gpu model
         if (this->test_model.device.compare("cuda") == 0) {
@@ -539,18 +536,18 @@ void CrossValidator::validate_backward(std::vector<float> &y_batch,
         }
     }
 
-    // Parameter update for input layer
-    if (test_model.param_update) {
-        test_model.layers[0]->param_backward(*test_model.layers[0]->bwd_states,
-                                             *test_input_delta_z_buffer,
-                                             *test_temp_states);
-    }
-
-    // State update for input layer
-    if (test_model.input_state_update) {
-        test_model.layers[0]->state_backward(
-            *test_model.layers[0]->bwd_states, *test_input_delta_z_buffer,
-            *test_output_delta_z_buffer, *test_temp_states);
-    }
+    // // Parameter update for input layer
+    // if (test_model.param_update) {
+    //     test_model.layers[0]->param_backward(*test_model.layers[0]->bwd_states,
+    //                                          *test_input_delta_z_buffer,
+    //                                          *test_temp_states);
+    // }
+
+    // // State update for input layer
+    // if (test_model.input_state_update) {
+    //     test_model.layers[0]->state_backward(
+    //         *test_model.layers[0]->bwd_states, *test_input_delta_z_buffer,
+    //         *test_output_delta_z_buffer, *test_temp_states);
+    // }
 }
 #endif
\ No newline at end of file
diff --git a/src/linear_layer.cpp b/src/linear_layer.cpp
index e7dba322a..9e634c414 100644
--- a/src/linear_layer.cpp
+++ b/src/linear_layer.cpp
@@ -564,78 +564,6 @@ void Linear::forward(BaseHiddenStates &input_states,
     }
 }
 
-void Linear::state_backward(BaseBackwardStates &next_bwd_states,
-                            BaseDeltaStates &input_delta_states,
-                            BaseDeltaStates &output_delta_states,
-                            BaseTempStates &temp_states)
-/*
- */
-{
-    // Initialization
-    int batch_size = input_delta_states.block_size;
-
-    // Compute inovation vector
-    if (this->num_threads > 1) {
-        linear_bwd_fc_delta_z_mp(
-            this->mu_w, next_bwd_states.jcb, input_delta_states.delta_mu,
-            input_delta_states.delta_var, this->input_size, this->output_size,
-            batch_size, this->num_threads, output_delta_states.delta_mu,
-            output_delta_states.delta_var);
-    } else {
-        int start_chunk = 0;
-        int end_chunk = batch_size * this->input_size;
-        linear_bwd_fc_delta_z(
-            this->mu_w, next_bwd_states.jcb, input_delta_states.delta_mu,
-            input_delta_states.delta_var, this->input_size, this->output_size,
-            batch_size, start_chunk, end_chunk, output_delta_states.delta_mu,
-            output_delta_states.delta_var);
-    }
-}
-
-void Linear::param_backward(BaseBackwardStates &next_bwd_states,
-                            BaseDeltaStates &delta_states,
-                            BaseTempStates &temp_states)
-/*
-...
-
-Args:
-    mu_a: Mean of activations from the next layer
- */
-{
-    // Initialization
-    int batch_size = delta_states.block_size;
-
-    // Update values for weights & biases
-    if (this->num_threads > 1) {
-        linear_bwd_fc_delta_w_mp(
-            this->var_w, next_bwd_states.mu_a, delta_states.delta_mu,
-            delta_states.delta_var, this->input_size, this->output_size,
-            batch_size, this->num_threads, this->delta_mu_w, this->delta_var_w);
-
-        if (this->bias) {
-            linear_bwd_fc_delta_b_mp(this->var_b, delta_states.delta_mu,
-                                     delta_states.delta_var, this->output_size,
-                                     batch_size, this->num_threads,
-                                     this->delta_mu_b, this->delta_var_b);
-        }
-    } else {
-        int start_chunk = 0;
-        int end_chunk = this->input_size * this->output_size;
-        linear_bwd_fc_delta_w(this->var_w, next_bwd_states.mu_a,
-                              delta_states.delta_mu, delta_states.delta_var,
-                              this->input_size, this->output_size, batch_size,
-                              start_chunk, end_chunk, this->delta_mu_w,
-                              this->delta_var_w);
-
-        if (this->bias) {
-            linear_bwd_fc_delta_b(this->var_b, delta_states.delta_mu,
-                                  delta_states.delta_var, this->output_size,
-                                  batch_size, start_chunk, this->output_size,
-                                  this->delta_mu_b, this->delta_var_b);
-        }
-    }
-}
-
 void Linear::backward(BaseDeltaStates &input_delta_states,
                       BaseDeltaStates &output_delta_states,
                       BaseTempStates &temp_states, bool state_udapte)
diff --git a/src/linear_layer_cuda.cu b/src/linear_layer_cuda.cu
index 8b38cce90..a5636522b 100644
--- a/src/linear_layer_cuda.cu
+++ b/src/linear_layer_cuda.cu
@@ -10,6 +10,9 @@
 #include "../include/linear_layer.h"
 #include "../include/linear_layer_cuda.cuh"
 
+#define WARP_SIZE 32U
+#define BANK_SIZE 32U
+
 __global__ void linear_fwd_mean_var(float const *mu_w, float const *var_w,
                                     float const *mu_b, float const *var_b,
                                     const float *mu_a, const float *var_a,
@@ -46,6 +49,325 @@ __global__ void linear_fwd_mean_var(float const *mu_w, float const *var_w,
     }
 }
 
+template <typename T, size_t BLOCK_TILE, size_t BLOCK_TILE_K, size_t THREADS>
+__global__ void linear_fwd_mean_var_v1(const T *mu_w, const T *var_w,
+                                       const T *mu_b, const T *var_b,
+                                       const T *mu_a, const T *var_a,
+                                       size_t input_size, size_t output_size,
+                                       int batch_size, bool bias, T *mu_z,
+                                       T *var_z)
+/*
+ */
+{
+    __shared__ T smem_mu_a[BLOCK_TILE][BLOCK_TILE_K];
+    __shared__ T smem_var_a[BLOCK_TILE][BLOCK_TILE_K];
+    __shared__ T smem_mu_w[BLOCK_TILE_K][BLOCK_TILE];
+    __shared__ T smem_var_w[BLOCK_TILE_K][BLOCK_TILE];
+
+    const size_t row = blockIdx.y * blockDim.y + threadIdx.y;
+    const size_t col = blockIdx.x * blockDim.x + threadIdx.x;
+    const size_t thread_linear_idx = threadIdx.y * blockDim.x + threadIdx.x;
+    unsigned int num_tiles = (input_size + BLOCK_TILE_K - 1) / BLOCK_TILE_K;
+    constexpr int NUM_LOADS =
+        (BLOCK_TILE * BLOCK_TILE_K + THREADS - 1) / THREADS;
+
+    T sum_mu = static_cast<T>(0);
+    T sum_var = static_cast<T>(0);
+
+    for (int phase = 0; phase < num_tiles; phase++) {
+        for (int l_i = 0; l_i < NUM_LOADS; l_i++) {
+            size_t a_ty = (thread_linear_idx + l_i * THREADS) / BLOCK_TILE_K;
+            size_t a_tx = (thread_linear_idx + l_i * THREADS) % BLOCK_TILE_K;
+
+            // input matrix to shared mem (batch_size x input_size)
+            size_t a_row_idx = blockIdx.y * BLOCK_TILE + a_ty;
+            size_t a_col_idx = phase * BLOCK_TILE_K + a_tx;
+            if (a_row_idx < batch_size && a_col_idx < input_size) {
+                smem_mu_a[a_ty][a_tx] =
+                    mu_a[a_row_idx * input_size + a_col_idx];
+                smem_var_a[a_ty][a_tx] =
+                    var_a[a_row_idx * input_size + a_col_idx];
+            } else {
+                smem_mu_a[a_ty][a_tx] = static_cast<T>(0);
+                smem_var_a[a_ty][a_tx] = static_cast<T>(0);
+            }
+
+            // weight matrix to shared mem (output_size x input_size)
+            size_t w_ty = (thread_linear_idx + l_i * THREADS) / BLOCK_TILE;
+            size_t w_tx = (thread_linear_idx + l_i * THREADS) % BLOCK_TILE;
+
+            size_t w_col_idx = phase * BLOCK_TILE_K + w_ty;
+            size_t w_row_idx = blockIdx.x * BLOCK_TILE + w_tx;
+
+            if (w_row_idx < output_size && w_col_idx < input_size) {
+                smem_mu_w[w_ty][w_tx] =
+                    mu_w[w_row_idx * input_size + w_col_idx];
+                smem_var_w[w_ty][w_tx] =
+                    var_w[w_row_idx * input_size + w_col_idx];
+            } else {
+                smem_mu_w[w_ty][w_tx] = static_cast<T>(0);
+                smem_var_w[w_ty][w_tx] = static_cast<T>(0);
+            }
+        }
+        __syncthreads();
+#pragma unroll
+        for (int i = 0; i < BLOCK_TILE_K; i++) {
+            T mu_w_tmp = smem_mu_w[i][threadIdx.x];
+            T var_w_tmp = smem_var_w[i][threadIdx.x];
+            T mu_a_tmp = smem_mu_a[threadIdx.y][i];
+            T var_a_tmp = smem_var_a[threadIdx.y][i];
+
+            sum_mu += mu_w_tmp * mu_a_tmp;
+            sum_var += (mu_w_tmp * mu_w_tmp + var_w_tmp) * var_a_tmp +
+                       var_w_tmp * mu_a_tmp * mu_a_tmp;
+        }
+        __syncthreads();
+    }
+    if (row < batch_size && col < output_size) {
+        mu_z[row * output_size + col] = bias ? sum_mu + mu_b[col] : sum_mu;
+        var_z[row * output_size + col] = bias ? sum_var + var_b[col] : sum_var;
+    }
+}
+
+template <typename T, size_t BLOCK_TILE, size_t BLOCK_TILE_K,
+          size_t THREAD_TILE, size_t THREADS, size_t SMEM_PADDING>
+__global__ void linear_fwd_mean_var_v2(const T *mu_w, const T *var_w,
+                                       const T *mu_b, const T *var_b,
+                                       const T *mu_a, const T *var_a,
+                                       size_t input_size, size_t output_size,
+                                       int batch_size, bool bias, T *mu_z,
+                                       T *var_z)
+/*
+ */
+{
+    __shared__ T smem_mu_a[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_var_a[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_mu_w[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_var_w[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+
+    const size_t tx = threadIdx.x;
+    const size_t ty = threadIdx.y;
+    const size_t thread_linear_idx = ty * blockDim.x + tx;
+
+    unsigned int num_tiles = (input_size + BLOCK_TILE_K - 1) / BLOCK_TILE_K;
+    constexpr int num_loads =
+        (BLOCK_TILE * BLOCK_TILE_K + THREADS - 1) / THREADS;
+
+    T tmp_mu[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_var[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T mu_w_val[THREAD_TILE] = {static_cast<T>(0)};
+    T var_w_val[THREAD_TILE] = {static_cast<T>(0)};
+    T mu_a_val[THREAD_TILE] = {static_cast<T>(0)};
+    T var_a_val[THREAD_TILE] = {static_cast<T>(0)};
+
+    for (size_t phase = 0; phase < num_tiles; phase++) {
+        for (int l_i = 0; l_i < num_loads; l_i++) {
+            // input matrix to shared mem (batch_size x input_size)
+            size_t a_ty = (thread_linear_idx + l_i * THREADS) / BLOCK_TILE_K;
+            size_t a_tx = (thread_linear_idx + l_i * THREADS) % BLOCK_TILE_K;
+
+            size_t a_row_idx = blockIdx.y * BLOCK_TILE + a_ty;
+            size_t a_col_idx = phase * BLOCK_TILE_K + a_tx;
+
+            if (a_row_idx < batch_size && a_col_idx < input_size) {
+                smem_mu_a[a_tx][a_ty] =
+                    mu_a[a_row_idx * input_size + a_col_idx];
+
+                smem_var_a[a_tx][a_ty] =
+                    var_a[a_row_idx * input_size + a_col_idx];
+            } else {
+                smem_mu_a[a_tx][a_ty] = static_cast<T>(0);
+                smem_var_a[a_tx][a_ty] = static_cast<T>(0);
+            }
+
+            // weight matrix to shared mem (output_size x input_size)
+            size_t w_ty = (thread_linear_idx + l_i * THREADS) / BLOCK_TILE;
+            size_t w_tx = (thread_linear_idx + l_i * THREADS) % BLOCK_TILE;
+
+            size_t w_col_idx = phase * BLOCK_TILE_K + w_ty;
+            size_t w_row_idx = blockIdx.x * BLOCK_TILE + w_tx;
+
+            if (w_row_idx < output_size && w_col_idx < input_size) {
+                smem_mu_w[w_ty][w_tx] =
+                    mu_w[w_row_idx * input_size + w_col_idx];
+                smem_var_w[w_ty][w_tx] =
+                    var_w[w_row_idx * input_size + w_col_idx];
+            } else {
+                smem_mu_w[w_ty][w_tx] = static_cast<T>(0);
+                smem_var_w[w_ty][w_tx] = static_cast<T>(0);
+            }
+        }
+        __syncthreads();
+#pragma unroll
+        for (size_t i = 0; i < BLOCK_TILE_K; i++) {
+#pragma unroll
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                mu_a_val[j] = smem_mu_a[i][ty * THREAD_TILE + j];
+                var_a_val[j] = smem_var_a[i][ty * THREAD_TILE + j];
+            }
+#pragma unroll
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                mu_w_val[j] = smem_mu_w[i][tx * THREAD_TILE + j];
+                var_w_val[j] = smem_var_w[i][tx * THREAD_TILE + j];
+            }
+#pragma unroll
+            for (size_t t = 0; t < THREAD_TILE; t++) {
+                for (size_t j = 0; j < THREAD_TILE; j++) {
+                    tmp_mu[t][j] += mu_w_val[j] * mu_a_val[t];
+                    tmp_var[t][j] +=
+                        (mu_w_val[j] * mu_w_val[j] + var_w_val[j]) *
+                            var_a_val[t] +
+                        var_w_val[j] * mu_a_val[t] * mu_a_val[t];
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+#pragma unroll
+    for (size_t t = 0; t < THREAD_TILE; t++) {
+        size_t row = blockIdx.y * BLOCK_TILE + ty * THREAD_TILE + t;
+#pragma unroll
+        for (size_t j = 0; j < THREAD_TILE; j++) {
+            size_t col = blockIdx.x * BLOCK_TILE + tx * THREAD_TILE + j;
+            if (row < batch_size && col < output_size) {
+                if (bias) {
+                    mu_z[row * output_size + col] = tmp_mu[t][j] + mu_b[col];
+                    var_z[row * output_size + col] = tmp_var[t][j] + var_b[col];
+                } else {
+                    mu_z[row * output_size + col] = tmp_mu[t][j];
+                    var_z[row * output_size + col] = tmp_var[t][j];
+                }
+            }
+        }
+    }
+}
+
+template <typename T, size_t BLOCK_TILE, size_t BLOCK_TILE_K,
+          size_t THREAD_TILE, size_t THREADS, size_t WARP_TILE_X,
+          size_t WARP_TILE_Y, size_t SMEM_PADDING>
+__global__ void linear_fwd_mean_var_v3(const T *mu_w, const T *var_w,
+                                       const T *mu_b, const T *var_b,
+                                       const T *mu_a, const T *var_a,
+                                       size_t input_size, size_t output_size,
+                                       int batch_size, bool bias, T *mu_z,
+                                       T *var_z)
+/*
+ */
+{
+    __shared__ T smem_mu_a[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_var_a[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_mu_w[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_var_w[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+
+    // Thread block
+    const size_t thread_linear_idx = threadIdx.y * blockDim.x + threadIdx.x;
+    unsigned int num_tiles = (input_size + BLOCK_TILE_K - 1) / BLOCK_TILE_K;
+    constexpr unsigned int NUM_LOADS =
+        (BLOCK_TILE * BLOCK_TILE_K + THREADS - 1) / THREADS;
+
+    // Warp
+    constexpr unsigned int WARPS_X = BLOCK_TILE / WARP_TILE_X;
+    constexpr unsigned int THREADS_PER_WARP_X = WARP_TILE_X / THREAD_TILE;
+    const size_t warp_id = thread_linear_idx / WARP_SIZE;
+    const size_t lane_id = thread_linear_idx % WARP_SIZE;
+    const size_t warp_row = warp_id / WARPS_X;
+    const size_t warp_col = warp_id % WARPS_X;
+    const size_t thread_row_in_warp = lane_id / THREADS_PER_WARP_X;
+    const size_t thread_col_in_warp = lane_id % THREADS_PER_WARP_X;
+
+    T tmp_mu[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_var[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T mu_w_val[THREAD_TILE] = {static_cast<T>(0)};
+    T var_w_val[THREAD_TILE] = {static_cast<T>(0)};
+    T mu_a_val[THREAD_TILE] = {static_cast<T>(0)};
+    T var_a_val[THREAD_TILE] = {static_cast<T>(0)};
+
+    for (size_t phase = 0; phase < num_tiles; phase++) {
+#pragma unroll
+        for (int l_i = 0; l_i < NUM_LOADS; l_i++) {
+            // input matrix to shared mem (batch_size x input_size)
+            size_t thread_load_idx = thread_linear_idx + l_i * THREADS;
+            size_t a_ty = thread_load_idx / BLOCK_TILE_K;
+            size_t a_tx = thread_load_idx % BLOCK_TILE_K;
+
+            size_t a_row = blockIdx.y * BLOCK_TILE + a_ty;
+            size_t a_col = phase * BLOCK_TILE_K + a_tx;
+
+            if (a_row < batch_size && a_col < input_size) {
+                smem_mu_a[a_tx][a_ty] = mu_a[a_row * input_size + a_col];
+                smem_var_a[a_tx][a_ty] = var_a[a_row * input_size + a_col];
+            } else {
+                smem_mu_a[a_tx][a_ty] = static_cast<T>(0);
+                smem_var_a[a_tx][a_ty] = static_cast<T>(0);
+            }
+
+            // weight matrix to shared mem (output_size x input_size)
+            size_t w_ty = thread_load_idx / BLOCK_TILE;
+            size_t w_tx = thread_load_idx % BLOCK_TILE;
+
+            size_t w_row = blockIdx.x * BLOCK_TILE + w_tx;
+            size_t w_col = phase * BLOCK_TILE_K + w_ty;
+
+            if (w_row < output_size && w_col < input_size) {
+                smem_mu_w[w_ty][w_tx] = mu_w[w_row * input_size + w_col];
+                smem_var_w[w_ty][w_tx] = var_w[w_row * input_size + w_col];
+            } else {
+                smem_mu_w[w_ty][w_tx] = static_cast<T>(0);
+                smem_var_w[w_ty][w_tx] = static_cast<T>(0);
+            }
+        }
+        __syncthreads();
+        for (size_t i = 0; i < BLOCK_TILE_K; i++) {
+#pragma unroll
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                size_t idx = warp_row * WARP_TILE_Y +
+                             thread_row_in_warp * THREAD_TILE + j;
+                mu_a_val[j] = smem_mu_a[i][idx];
+                var_a_val[j] = smem_var_a[i][idx];
+            }
+#pragma unroll
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                size_t idx = warp_col * WARP_TILE_X +
+                             thread_col_in_warp * THREAD_TILE + j;
+                mu_w_val[j] = smem_mu_w[i][idx];
+                var_w_val[j] = smem_var_w[i][idx];
+            }
+#pragma unroll
+            for (size_t t = 0; t < THREAD_TILE; t++) {
+                for (size_t j = 0; j < THREAD_TILE; j++) {
+                    tmp_mu[t][j] += mu_w_val[j] * mu_a_val[t];
+                    tmp_var[t][j] +=
+                        (mu_w_val[j] * mu_w_val[j] + var_w_val[j]) *
+                            var_a_val[t] +
+                        var_w_val[j] * mu_a_val[t] * mu_a_val[t];
+                }
+            }
+        }
+        __syncwarp();
+    }
+    __syncthreads();
+
+    const size_t base_row = blockIdx.y * BLOCK_TILE + warp_row * WARP_TILE_Y +
+                            thread_row_in_warp * THREAD_TILE;
+    const size_t base_col = blockIdx.x * BLOCK_TILE + warp_col * WARP_TILE_X +
+                            thread_col_in_warp * THREAD_TILE;
+
+    for (size_t t = 0; t < THREAD_TILE; t++) {
+        size_t row = base_row + t;
+#pragma unroll
+        for (size_t j = 0; j < THREAD_TILE; j++) {
+            size_t col = base_col + j;
+            if (row < batch_size && col < output_size) {
+                mu_z[row * output_size + col] =
+                    bias ? tmp_mu[t][j] + mu_b[col] : tmp_mu[t][j];
+                var_z[row * output_size + col] =
+                    bias ? tmp_var[t][j] + var_b[col] : tmp_var[t][j];
+            }
+        }
+    }
+}
+
 __global__ void linear_fwd_full_cov(float const *mu_w, float const *var_a_f,
                                     size_t input_size, size_t output_size,
                                     int batch_size, float *var_z_fp)
@@ -143,6 +465,219 @@ __global__ void linear_bwd_delta_z(float const *mu_w, float const *jcb,
     }
 }
 
+template <typename T, size_t BLOCK_TILE, size_t BLOCK_TILE_K,
+          size_t THREAD_TILE, size_t THREADS, size_t SMEM_PADDING>
+__global__ void linear_bwd_delta_z_v2(const T *mu_w, const T *jcb,
+                                      const T *delta_mu_out,
+                                      const T *delta_var_out, size_t input_size,
+                                      size_t output_size, int batch_size,
+                                      T *delta_mu_in, T *delta_var_in)
+/*
+ */
+{
+    __shared__ T smem_delta_mu[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_delta_var[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_mu_w[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+
+    const size_t tx = threadIdx.x;
+    const size_t ty = threadIdx.y;
+    const size_t thread_linear_idx = ty * blockDim.x + tx;
+
+    unsigned int num_tiles = (output_size + BLOCK_TILE_K - 1) / BLOCK_TILE_K;
+
+    constexpr int NUM_LOADS =
+        (BLOCK_TILE * BLOCK_TILE_K + THREADS - 1) / THREADS;
+
+    T tmp_mu[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_var[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T mu_w_val[THREAD_TILE] = {static_cast<T>(0)};
+    T delta_mu_val[THREAD_TILE] = {static_cast<T>(0)};
+    T delta_var_val[THREAD_TILE] = {static_cast<T>(0)};
+
+    for (size_t phase = 0; phase < num_tiles; phase++) {
+        for (int l_i = 0; l_i < NUM_LOADS; l_i++) {
+            const int thread_load_idx = thread_linear_idx + l_i * THREADS;
+
+            // delta matrix
+            size_t delta_ty = thread_load_idx / BLOCK_TILE_K;
+            size_t delta_tx = thread_load_idx % BLOCK_TILE_K;
+
+            size_t delta_row = blockIdx.y * BLOCK_TILE + delta_ty;
+            size_t delta_col = phase * BLOCK_TILE_K + delta_tx;
+
+            if (delta_row < batch_size && delta_col < output_size) {
+                smem_delta_mu[delta_tx][delta_ty] =
+                    delta_mu_out[delta_row * output_size + delta_col];
+                smem_delta_var[delta_tx][delta_ty] =
+                    delta_var_out[delta_row * output_size + delta_col];
+            } else {
+                smem_delta_mu[delta_tx][delta_ty] = static_cast<T>(0);
+                smem_delta_var[delta_tx][delta_ty] = static_cast<T>(0);
+            }
+
+            // weight matrix
+            size_t w_ty = thread_load_idx / BLOCK_TILE;
+            size_t w_tx = thread_load_idx % BLOCK_TILE;
+            size_t w_row = phase * BLOCK_TILE_K + w_ty;
+            size_t w_col = blockIdx.x * BLOCK_TILE + w_tx;
+
+            if (w_row < output_size && w_col < input_size) {
+                smem_mu_w[w_ty][w_tx] = mu_w[w_row * input_size + w_col];
+            } else {
+                smem_mu_w[w_ty][w_tx] = static_cast<T>(0);
+            }
+        }
+        __syncthreads();
+        for (size_t i = 0; i < BLOCK_TILE_K; i++) {
+#pragma unroll
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                delta_mu_val[j] = smem_delta_mu[i][ty * THREAD_TILE + j];
+                delta_var_val[j] = smem_delta_var[i][ty * THREAD_TILE + j];
+            }
+#pragma unroll
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                mu_w_val[j] = smem_mu_w[i][tx * THREAD_TILE + j];
+            }
+
+            for (size_t t = 0; t < THREAD_TILE; t++) {
+#pragma unroll
+                for (size_t j = 0; j < THREAD_TILE; j++) {
+                    tmp_mu[t][j] += mu_w_val[j] * delta_mu_val[t];
+                    tmp_var[t][j] +=
+                        mu_w_val[j] * delta_var_val[t] * mu_w_val[j];
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    for (size_t t = 0; t < THREAD_TILE; t++) {
+        int row = blockIdx.y * BLOCK_TILE + ty * THREAD_TILE + t;
+#pragma unroll
+        for (size_t j = 0; j < THREAD_TILE; j++) {
+            int col = blockIdx.x * BLOCK_TILE + tx * THREAD_TILE + j;
+            if (row < batch_size && col < input_size) {
+                T jcb_val = jcb[row * input_size + col];
+
+                delta_mu_in[row * input_size + col] = tmp_mu[t][j] * jcb_val;
+                delta_var_in[row * input_size + col] =
+                    tmp_var[t][j] * jcb_val * jcb_val;
+            }
+        }
+    }
+}
+template <typename T, size_t BLOCK_TILE, size_t BLOCK_TILE_K,
+          size_t THREAD_TILE, size_t THREADS, size_t WARP_TILE_X,
+          size_t WARP_TILE_Y, size_t SMEM_PADDING>
+__global__ void linear_bwd_delta_z_v3(const T *mu_w, const T *jcb,
+                                      const T *delta_mu_out,
+                                      const T *delta_var_out, size_t input_size,
+                                      size_t output_size, int batch_size,
+                                      T *delta_mu_in, T *delta_var_in)
+/*
+ */
+{
+    __shared__ T smem_delta_mu[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_delta_var[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_mu_w[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+
+    const size_t thread_linear_idx = threadIdx.y * blockDim.x + threadIdx.x;
+    unsigned int num_tiles = (output_size + BLOCK_TILE_K - 1) / BLOCK_TILE_K;
+    constexpr unsigned int WARPS_X = BLOCK_TILE / WARP_TILE_X;
+    constexpr unsigned int THREADS_PER_WARP_X = WARP_TILE_X / THREAD_TILE;
+    constexpr int NUM_LOADS =
+        (BLOCK_TILE * BLOCK_TILE_K + THREADS - 1) / THREADS;
+
+    const size_t warp_id = thread_linear_idx / WARP_SIZE;
+    const size_t lane_id = thread_linear_idx % WARP_SIZE;
+    const size_t warp_row = warp_id / WARPS_X;
+    const size_t warp_col = warp_id % WARPS_X;
+    const size_t thread_row_in_warp = lane_id / THREADS_PER_WARP_X;
+    const size_t thread_col_in_warp = lane_id % THREADS_PER_WARP_X;
+
+    T tmp_mu[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_var[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_delta_mu[THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_delta_var[THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_mu_w[THREAD_TILE] = {static_cast<T>(0)};
+
+    for (size_t phase = 0; phase < num_tiles; phase++) {
+        for (int l_i = 0; l_i < NUM_LOADS; l_i++) {
+            const int thread_load_idx = thread_linear_idx + l_i * THREADS;
+
+            // delta matrix
+            size_t delta_ty = thread_load_idx / BLOCK_TILE_K;
+            size_t delta_tx = thread_load_idx % BLOCK_TILE_K;
+
+            size_t delta_row = blockIdx.y * BLOCK_TILE + delta_ty;
+            size_t delta_col = phase * BLOCK_TILE_K + delta_tx;
+
+            if (delta_row < batch_size && delta_col < output_size) {
+                smem_delta_mu[delta_tx][delta_ty] =
+                    delta_mu_out[delta_row * output_size + delta_col];
+                smem_delta_var[delta_tx][delta_ty] =
+                    delta_var_out[delta_row * output_size + delta_col];
+            } else {
+                smem_delta_mu[delta_tx][delta_ty] = static_cast<T>(0);
+                smem_delta_var[delta_tx][delta_ty] = static_cast<T>(0);
+            }
+
+            // weight matrix
+            size_t w_ty = thread_load_idx / BLOCK_TILE;
+            size_t w_tx = thread_load_idx % BLOCK_TILE;
+            size_t w_row = phase * BLOCK_TILE_K + w_ty;
+            size_t w_col = blockIdx.x * BLOCK_TILE + w_tx;
+
+            if (w_row < output_size && w_col < input_size) {
+                smem_mu_w[w_ty][w_tx] = mu_w[w_row * input_size + w_col];
+            } else {
+                smem_mu_w[w_ty][w_tx] = static_cast<T>(0);
+            }
+        }
+        __syncthreads();
+
+        for (size_t i = 0; i < BLOCK_TILE_K; i++) {
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                size_t idx = warp_row * WARP_TILE_Y +
+                             thread_row_in_warp * THREAD_TILE + j;
+                tmp_delta_mu[j] = smem_delta_mu[i][idx];
+                tmp_delta_var[j] = smem_delta_var[i][idx];
+            }
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                size_t idx = warp_col * WARP_TILE_X +
+                             thread_col_in_warp * THREAD_TILE + j;
+                tmp_mu_w[j] = smem_mu_w[i][idx];
+            }
+
+            for (size_t t = 0; t < THREAD_TILE; t++) {
+                for (size_t j = 0; j < THREAD_TILE; j++) {
+                    tmp_mu[t][j] += tmp_mu_w[j] * tmp_delta_mu[t];
+                    tmp_var[t][j] +=
+                        tmp_mu_w[j] * tmp_delta_var[t] * tmp_mu_w[j];
+                }
+            }
+        }
+        __syncwarp();
+    }
+    __syncthreads();
+
+    for (size_t t = 0; t < THREAD_TILE; t++) {
+        int row = blockIdx.y * BLOCK_TILE + warp_row * WARP_TILE_Y +
+                  thread_row_in_warp * THREAD_TILE + t;
+#pragma unroll
+        for (size_t j = 0; j < THREAD_TILE; j++) {
+            int col = blockIdx.x * BLOCK_TILE + warp_col * WARP_TILE_X +
+                      thread_col_in_warp * THREAD_TILE + j;
+            if (row < batch_size && col < input_size) {
+                T jcb_val = jcb[row * input_size + col];
+                delta_mu_in[row * input_size + col] = tmp_mu[t][j] * jcb_val;
+                delta_var_in[row * input_size + col] =
+                    tmp_var[t][j] * jcb_val * jcb_val;
+            }
+        }
+    }
+}
+
 __global__ void linear_bwd_delta_w(float const *var_w, float const *mu_a,
                                    float const *delta_mu_out,
                                    float const *delta_var_out,
@@ -174,6 +709,223 @@ __global__ void linear_bwd_delta_w(float const *var_w, float const *mu_a,
     }
 }
 
+template <typename T, size_t BLOCK_TILE, size_t BLOCK_TILE_K,
+          size_t THREAD_TILE, size_t THREADS, size_t SMEM_PADDING>
+__global__ void linear_bwd_delta_w_v2(const T *var_w, const T *mu_a,
+                                      const T *delta_mu_out,
+                                      const T *delta_var_out, size_t input_size,
+                                      size_t output_size, int batch_size,
+                                      T *delta_mu_w, T *delta_var_w)
+/**/
+{
+    __shared__ T smem_delta_mu[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_delta_var[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_mu_a[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+
+    const size_t ty = threadIdx.y;
+    const size_t tx = threadIdx.x;
+    const size_t thread_linear_idx = ty * blockDim.x + tx;
+
+    constexpr unsigned int NUM_LOADS =
+        (BLOCK_TILE * BLOCK_TILE_K + THREADS - 1) / THREADS;
+    unsigned int num_tiles = (batch_size + BLOCK_TILE_K - 1) / BLOCK_TILE_K;
+
+    T tmp_mu[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_var[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_delta_mu[THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_delta_var[THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_mu_a[THREAD_TILE] = {static_cast<T>(0)};
+
+    for (size_t phase = 0; phase < num_tiles; phase++) {
+        for (size_t l_i = 0; l_i < NUM_LOADS; l_i++) {
+            const size_t thread_load_idx = thread_linear_idx + l_i * THREADS;
+
+            // Update quantity matrix
+            size_t delta_ty = thread_load_idx / BLOCK_TILE_K;
+            size_t delta_tx = thread_load_idx % BLOCK_TILE_K;
+            size_t delta_col = blockIdx.y * BLOCK_TILE + delta_ty;
+            size_t delta_row = phase * BLOCK_TILE_K + delta_tx;
+
+            smem_delta_mu[delta_tx][delta_ty] =
+                (delta_row < batch_size && delta_col < output_size)
+                    ? delta_mu_out[delta_row * output_size + delta_col]
+                    : static_cast<T>(0);
+
+            smem_delta_var[delta_tx][delta_ty] =
+                (delta_row < batch_size && delta_col < output_size)
+                    ? delta_var_out[delta_row * output_size + delta_col]
+                    : static_cast<T>(0);
+
+            // activation matrix
+            size_t a_ty = thread_load_idx / BLOCK_TILE;
+            size_t a_tx = thread_load_idx & BLOCK_TILE;
+            size_t a_col = blockIdx.x * BLOCK_TILE + a_tx;
+            size_t a_row = phase * BLOCK_TILE_K + a_ty;
+
+            smem_mu_a[a_ty][a_tx] = (a_row < batch_size && a_col < input_size)
+                                        ? smem_mu_a[a_ty][a_tx] =
+                                              mu_a[a_row * input_size + a_col]
+                                        : static_cast<T>(0);
+        }
+        __syncthreads();
+
+        for (size_t i = 0; i < BLOCK_TILE_K; i++) {
+#pragma unroll
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                tmp_delta_mu[j] = smem_delta_mu[i][ty * THREAD_TILE + j];
+                tmp_delta_var[j] = smem_delta_var[i][ty * THREAD_TILE + j];
+            }
+#pragma unroll
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                tmp_mu_a[j] = smem_mu_a[i][tx * THREAD_TILE + j];
+            }
+
+            for (size_t t = 0; t < THREAD_TILE; t++) {
+#pragma unroll
+                for (size_t j = 0; j < THREAD_TILE; j++) {
+                    tmp_mu[t][j] += tmp_mu_a[j] * tmp_delta_mu[t];
+                    tmp_var[t][j] +=
+                        tmp_mu_a[j] * tmp_mu_a[j] * tmp_delta_var[t];
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    for (size_t t = 0; t < THREAD_TILE; t++) {
+        size_t row = blockIdx.y * BLOCK_TILE + ty * THREAD_TILE + t;
+#pragma unroll
+        for (size_t j = 0; j < THREAD_TILE; j++) {
+            size_t col = blockIdx.x * BLOCK_TILE + tx * THREAD_TILE + j;
+            if (row < output_size && col < input_size) {
+                T tmp_var_w = var_w[row * input_size + col];
+
+                delta_mu_w[row * input_size + col] = tmp_mu[t][j] * tmp_var_w;
+                delta_var_w[row * input_size + col] =
+                    tmp_var[t][j] * tmp_var_w * tmp_var_w;
+            }
+        }
+    }
+}
+
+template <typename T, size_t BLOCK_TILE, size_t BLOCK_TILE_K,
+          size_t THREAD_TILE, size_t THREADS, size_t WARP_TILE_X,
+          size_t WARP_TILE_Y, size_t SMEM_PADDING>
+__global__ void linear_bwd_delta_w_v3(const T *var_w, const T *mu_a,
+                                      const T *delta_mu_out,
+                                      const T *delta_var_out, size_t input_size,
+                                      size_t output_size, int batch_size,
+                                      T *delta_mu_w, T *delta_var_w)
+/*
+ */
+{
+    __shared__ T smem_delta_mu[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_delta_var[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+    __shared__ T smem_mu_a[BLOCK_TILE_K][BLOCK_TILE + SMEM_PADDING];
+
+    constexpr unsigned int WAPRS_X = BLOCK_TILE / WARP_TILE_X;
+    constexpr unsigned int THREADS_PER_WARP_X = WARP_TILE_X / THREAD_TILE;
+    constexpr unsigned int NUM_LOADS =
+        (BLOCK_TILE * BLOCK_TILE_K + THREADS - 1) / THREADS;
+
+    const size_t thread_linear_idx = threadIdx.y * blockDim.x + threadIdx.x;
+    const size_t warp_id = thread_linear_idx / WARP_SIZE;
+    const size_t lane_id = thread_linear_idx % WARP_SIZE;
+    const size_t warp_row = warp_id / WAPRS_X;
+    const size_t warp_col = warp_id % WAPRS_X;
+    const size_t thread_row_per_warp = lane_id / THREADS_PER_WARP_X;
+    const size_t thread_col_per_warp = lane_id % THREADS_PER_WARP_X;
+
+    const unsigned int num_tiles =
+        (batch_size + BLOCK_TILE_K - 1) / BLOCK_TILE_K;
+
+    T tmp_mu[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_var[THREAD_TILE][THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_delta_mu[THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_delta_var[THREAD_TILE] = {static_cast<T>(0)};
+    T tmp_mu_a[THREAD_TILE] = {static_cast<T>(0)};
+
+    for (size_t phase = 0; phase < num_tiles; phase++) {
+        for (size_t l_i = 0; l_i < NUM_LOADS; l_i++) {
+            const size_t thread_load_idx = thread_linear_idx + l_i * THREADS;
+
+            // Update quantity matrix
+            size_t delta_ty = thread_load_idx / BLOCK_TILE_K;
+            size_t delta_tx = thread_load_idx % BLOCK_TILE_K;
+            size_t delta_col = blockIdx.y * BLOCK_TILE + delta_ty;
+            size_t delta_row = phase * BLOCK_TILE_K + delta_tx;
+
+            smem_delta_mu[delta_tx][delta_ty] =
+                (delta_row < batch_size && delta_col < output_size)
+                    ? delta_mu_out[delta_row * output_size + delta_col]
+                    : static_cast<T>(0);
+
+            smem_delta_var[delta_tx][delta_ty] =
+                (delta_row < batch_size && delta_col < output_size)
+                    ? delta_var_out[delta_row * output_size + delta_col]
+                    : static_cast<T>(0);
+
+            // activation matrix
+            size_t a_ty = thread_load_idx / BLOCK_TILE;
+            size_t a_tx = thread_load_idx % BLOCK_TILE;
+            size_t a_col = blockIdx.x * BLOCK_TILE + a_tx;
+            size_t a_row = phase * BLOCK_TILE_K + a_ty;
+
+            smem_mu_a[a_ty][a_tx] = (a_row < batch_size && a_col < input_size)
+                                        ? smem_mu_a[a_ty][a_tx] =
+                                              mu_a[a_row * input_size + a_col]
+                                        : static_cast<T>(0);
+        }
+        __syncthreads();
+
+        for (size_t i = 0; i < BLOCK_TILE_K; i++) {
+#pragma unroll
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                size_t idx = warp_row * WARP_TILE_Y +
+                             thread_row_per_warp * THREAD_TILE + j;
+                tmp_delta_mu[j] = smem_delta_mu[i][idx];
+                tmp_delta_var[j] = smem_delta_var[i][idx];
+            }
+#pragma unroll
+            for (size_t j = 0; j < THREAD_TILE; j++) {
+                size_t idx = warp_col * WARP_TILE_X +
+                             thread_col_per_warp * THREAD_TILE + j;
+                tmp_mu_a[j] = smem_mu_a[i][idx];
+            }
+
+            for (size_t t = 0; t < THREAD_TILE; t++) {
+#pragma unroll
+                for (size_t j = 0; j < THREAD_TILE; j++) {
+                    tmp_mu[t][j] += tmp_mu_a[j] * tmp_delta_mu[t];
+                    tmp_var[t][j] +=
+                        tmp_mu_a[j] * tmp_mu_a[j] * tmp_delta_var[t];
+                }
+            }
+        }
+        __syncwarp();
+    }
+    __syncthreads();
+    size_t row_base = blockIdx.y * BLOCK_TILE + warp_row * WARP_TILE_Y +
+                      thread_row_per_warp * THREAD_TILE;
+    size_t col_base = blockIdx.x * BLOCK_TILE + warp_col * WARP_TILE_X +
+                      thread_col_per_warp * THREAD_TILE;
+    for (size_t t = 0; t < THREAD_TILE; t++) {
+        size_t row = row_base + t;
+#pragma unroll
+        for (size_t j = 0; j < THREAD_TILE; j++) {
+            size_t col = col_base + j;
+
+            if (row < output_size && col < input_size) {
+                T tmp_var_w = var_w[row * input_size + col];
+
+                delta_mu_w[row * input_size + col] = tmp_mu[t][j] * tmp_var_w;
+                delta_var_w[row * input_size + col] =
+                    tmp_var[t][j] * tmp_var_w * tmp_var_w;
+            }
+        }
+    }
+}
+
 __global__ void linear_bwd_delta_b(float const *var_b,
                                    float const *delta_mu_out,
                                    float const *delta_var_out,
@@ -206,6 +958,128 @@ __global__ void linear_bwd_delta_b(float const *var_b,
 ////////////////////////////////////////////////////////////////////////////////
 // Fully Connected Layer
 ////////////////////////////////////////////////////////////////////////////////
+void linear_forward_cuda(HiddenStateCuda *&cu_input_states,
+                         HiddenStateCuda *&cu_output_states,
+                         const float *d_mu_w, const float *d_var_w,
+                         const float *d_mu_b, const float *d_var_b,
+                         size_t input_size, size_t output_size, int batch_size,
+                         bool bias) {
+    if (output_size > 128 && input_size > 1024) {
+        // TODO: remove hardcoded kernel config
+        constexpr unsigned int BLOCK_SIZE = 128U;
+        constexpr unsigned int THREAD_TILE = 8U;
+        constexpr unsigned int WARPS_X = 4U;
+        constexpr unsigned int WARPS_Y = 2U;
+        constexpr unsigned int BLOCK_TILE_K = 16U;
+        constexpr unsigned int WARP_TILE_X = BLOCK_SIZE / WARPS_X;
+        constexpr unsigned int WARP_TILE_Y = BLOCK_SIZE / WARPS_Y;
+        constexpr unsigned int THREADS_X =
+            WARPS_X * (WARP_TILE_X / THREAD_TILE);
+        constexpr unsigned int THREADS_Y =
+            WARPS_Y * (WARP_TILE_Y / THREAD_TILE);
+        constexpr unsigned int THREADS = THREADS_X * THREADS_Y;
+        constexpr unsigned int SMEM_PADDING = BANK_SIZE / THREADS_X;
+
+        dim3 block_dim(THREADS_X, THREADS_Y, 1U);
+        dim3 grid_dim(
+            (static_cast<unsigned int>(output_size) + BLOCK_SIZE - 1U) /
+                BLOCK_SIZE,
+            (static_cast<unsigned int>(batch_size) + BLOCK_SIZE - 1U) /
+                BLOCK_SIZE,
+            1U);
+
+        linear_fwd_mean_var_v3<float, BLOCK_SIZE, BLOCK_TILE_K, THREAD_TILE,
+                               THREADS, WARP_TILE_X, WARP_TILE_Y, SMEM_PADDING>
+            <<<grid_dim, block_dim>>>(
+                d_mu_w, d_var_w, d_mu_b, d_var_b, cu_input_states->d_mu_a,
+                cu_input_states->d_var_a, input_size, output_size, batch_size,
+                bias, cu_output_states->d_mu_a, cu_output_states->d_var_a);
+    } else {
+        constexpr unsigned int BLOCK_SIZE = 16U;
+        unsigned int grid_rows = (batch_size + BLOCK_SIZE - 1) / BLOCK_SIZE;
+        unsigned int grid_cols = (output_size + BLOCK_SIZE - 1) / BLOCK_SIZE;
+        dim3 block_dim(BLOCK_SIZE, BLOCK_SIZE, 1U);
+        dim3 grid_dim(grid_cols, grid_rows, 1U);
+        constexpr size_t THREADS = BLOCK_SIZE * BLOCK_SIZE;
+
+        linear_fwd_mean_var_v1<float, BLOCK_SIZE, BLOCK_SIZE, THREADS>
+            <<<grid_dim, block_dim>>>(
+                d_mu_w, d_var_w, d_mu_b, d_var_b, cu_input_states->d_mu_a,
+                cu_input_states->d_var_a, input_size, output_size, batch_size,
+                bias, cu_output_states->d_mu_a, cu_output_states->d_var_a);
+    }
+}
+
+void linear_state_backward_cuda(DeltaStateCuda *&cu_input_delta_states,
+                                DeltaStateCuda *&cu_output_delta_states,
+                                BackwardStateCuda *&cu_next_bwd_states,
+                                const float *d_mu_w, size_t input_size,
+                                size_t output_size, int batch_size)
+/*
+ */
+{
+    constexpr unsigned int BLOCK_SIZE = 128U;
+    constexpr unsigned int THREAD_TILE = 8U;
+    constexpr unsigned int WARPS_X = 4U;
+    constexpr unsigned int WARPS_Y = 2U;
+    constexpr unsigned int BLOCK_TILE_K = 16U;
+    constexpr unsigned int WARP_TILE_X = BLOCK_SIZE / WARPS_X;
+    constexpr unsigned int WARP_TILE_Y = BLOCK_SIZE / WARPS_Y;
+    constexpr unsigned int THREADS_X = WARPS_X * (WARP_TILE_X / THREAD_TILE);
+    constexpr unsigned int THREADS_Y = WARPS_Y * (WARP_TILE_Y / THREAD_TILE);
+    constexpr unsigned int THREADS = THREADS_X * THREADS_Y;
+    constexpr unsigned int SMEM_PADDING = BANK_SIZE / THREADS_X;
+
+    dim3 block_dim(THREADS_X, THREADS_Y, 1U);
+    dim3 grid_dim(
+        (static_cast<unsigned int>(input_size) + BLOCK_SIZE - 1U) / BLOCK_SIZE,
+        (static_cast<unsigned int>(batch_size) + BLOCK_SIZE - 1U) / BLOCK_SIZE,
+        1U);
+    linear_bwd_delta_z_v3<float, BLOCK_SIZE, BLOCK_TILE_K, THREAD_TILE, THREADS,
+                          WARP_TILE_X, WARP_TILE_Y, SMEM_PADDING>
+        <<<grid_dim, block_dim>>>(d_mu_w, cu_next_bwd_states->d_jcb,
+                                  cu_input_delta_states->d_delta_mu,
+                                  cu_input_delta_states->d_delta_var,
+                                  input_size, output_size, batch_size,
+                                  cu_output_delta_states->d_delta_mu,
+                                  cu_output_delta_states->d_delta_var);
+}
+
+void linear_weight_backward_cuda(DeltaStateCuda *&cu_input_delta_states,
+                                 DeltaStateCuda *&cu_output_delta_states,
+                                 BackwardStateCuda *&cu_next_bwd_states,
+                                 const float *d_var_w, size_t input_size,
+                                 size_t output_size, int batch_size,
+                                 float *d_delta_mu_w, float *d_delta_var_w)
+/*
+ */
+{
+    constexpr unsigned int BLOCK_SIZE = 128U;
+    constexpr unsigned int THREAD_TILE = 8U;
+    constexpr unsigned int WARPS_X = 4U;
+    constexpr unsigned int WARPS_Y = 2U;
+    constexpr unsigned int BLOCK_TILE_K = 16U;
+    constexpr unsigned int WARP_TILE_X = BLOCK_SIZE / WARPS_X;
+    constexpr unsigned int WARP_TILE_Y = BLOCK_SIZE / WARPS_Y;
+    constexpr unsigned int THREADS_X = WARPS_X * (WARP_TILE_X / THREAD_TILE);
+    constexpr unsigned int THREADS_Y = WARPS_Y * (WARP_TILE_Y / THREAD_TILE);
+    constexpr unsigned int THREADS = THREADS_X * THREADS_Y;
+    constexpr unsigned int SMEM_PADDING = BANK_SIZE / THREADS_X;
+
+    dim3 block_dim(THREADS_X, THREADS_Y, 1U);
+    dim3 grid_dim(
+        (static_cast<unsigned int>(input_size) + BLOCK_SIZE - 1U) / BLOCK_SIZE,
+        (static_cast<unsigned int>(output_size) + BLOCK_SIZE - 1U) / BLOCK_SIZE,
+        1U);
+
+    linear_bwd_delta_w_v3<float, BLOCK_SIZE, BLOCK_TILE_K, THREAD_TILE, THREADS,
+                          WARP_TILE_X, WARP_TILE_Y, SMEM_PADDING>
+        <<<grid_dim, block_dim>>>(d_var_w, cu_next_bwd_states->d_mu_a,
+                                  cu_input_delta_states->d_delta_mu,
+                                  cu_input_delta_states->d_delta_var,
+                                  input_size, output_size, batch_size,
+                                  d_delta_mu_w, d_delta_var_w);
+}
 
 LinearCuda::LinearCuda(size_t ip_size, size_t op_size, bool bias,
                        float gain_weight, float gain_bias, std::string method)
@@ -285,21 +1159,13 @@ void LinearCuda::forward(BaseHiddenStates &input_states,
     // *>(&temp_states);
 
     int batch_size = input_states.block_size;
-    int threads = this->num_cuda_threads;
-    this->set_cap_factor_udapte(batch_size);
 
-    // Forward pass
-    unsigned int grid_rows = (this->output_size + threads - 1) / threads;
-    unsigned int grid_cols = (batch_size + threads - 1) / threads;
-
-    dim3 grid_dim(grid_cols, grid_rows);
-    dim3 block_dim(threads, threads);
+    this->set_cap_factor_udapte(batch_size);
 
-    linear_fwd_mean_var<<<grid_dim, block_dim>>>(
-        this->d_mu_w, this->d_var_w, this->d_mu_b, this->d_var_b,
-        cu_input_states->d_mu_a, cu_input_states->d_var_a, this->input_size,
-        this->output_size, input_states.block_size, this->bias,
-        cu_output_states->d_mu_a, cu_output_states->d_var_a);
+    linear_forward_cuda(cu_input_states, cu_output_states, this->d_mu_w,
+                        this->d_var_w, this->d_mu_b, this->d_var_b,
+                        this->input_size, this->output_size,
+                        input_states.block_size, this->bias);
 
     // Update number of actual states.
     output_states.block_size = batch_size;
@@ -323,79 +1189,6 @@ void LinearCuda::forward(BaseHiddenStates &input_states,
     }
 }
 
-void LinearCuda::state_backward(BaseBackwardStates &next_bwd_states,
-                                BaseDeltaStates &input_delta_states,
-                                BaseDeltaStates &output_delta_states,
-                                BaseTempStates &temp_states)
-/*
- */
-{
-    // New poitner will point to the same memory location when casting
-    BackwardStateCuda *cu_next_bwd_states =
-        dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
-    DeltaStateCuda *cu_input_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&input_delta_states);
-    DeltaStateCuda *cu_output_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&output_delta_states);
-
-    // Initialization
-    int batch_size = input_delta_states.block_size;
-    int threads = this->num_cuda_threads;
-
-    // Compute inovation vector
-    unsigned int grid_row = (this->input_size + threads - 1) / threads;
-    unsigned int grid_col = (batch_size + threads - 1) / threads;
-
-    dim3 grid_dim(grid_col, grid_row);
-    dim3 block_dim(threads, threads);
-
-    linear_bwd_delta_z<<<grid_dim, block_dim>>>(
-        this->d_mu_w, cu_next_bwd_states->d_jcb,
-        cu_input_delta_states->d_delta_mu, cu_input_delta_states->d_delta_var,
-        this->input_size, this->output_size, batch_size,
-        cu_output_delta_states->d_delta_mu,
-        cu_output_delta_states->d_delta_var);
-}
-
-void LinearCuda::param_backward(BaseBackwardStates &next_bwd_states,
-                                BaseDeltaStates &delta_states,
-                                BaseTempStates &temp_states)
-/*
- */
-{
-    // New poitner will point to the same memory location when casting
-    BackwardStateCuda *cu_next_bwd_states =
-        dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
-    DeltaStateCuda *cu_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&delta_states);
-
-    // Initalization
-    int batch_size = delta_states.block_size;
-    int threads = this->num_cuda_threads;
-    dim3 block_dim(threads, threads);
-
-    // Updated values for weights
-    unsigned int grid_row_w = (this->input_size + threads - 1) / threads;
-    unsigned int grid_col_w = (this->output_size + threads - 1) / threads;
-    dim3 grid_dim_w(grid_col_w, grid_row_w);
-
-    linear_bwd_delta_w<<<grid_dim_w, block_dim>>>(
-        this->d_var_w, cu_next_bwd_states->d_mu_a, cu_delta_states->d_delta_mu,
-        cu_delta_states->d_delta_var, this->input_size, this->output_size,
-        batch_size, this->d_delta_mu_w, this->d_delta_var_w);
-
-    // Updated values for biases
-    if (this->bias) {
-        unsigned int grid_row_b = (this->output_size + threads - 1) / threads;
-        dim3 grid_dim_b(1, grid_row_b);
-
-        linear_bwd_delta_b<<<grid_dim_b, block_dim>>>(
-            this->d_var_b, cu_delta_states->d_delta_mu,
-            cu_delta_states->d_delta_var, this->input_size, this->output_size,
-            batch_size, this->d_delta_mu_b, this->d_delta_var_b);
-    }
-}
-
 void LinearCuda::backward(BaseDeltaStates &input_delta_states,
                           BaseDeltaStates &output_delta_states,
                           BaseTempStates &temp_states, bool state_udapte)
@@ -419,27 +1212,19 @@ void LinearCuda::backward(BaseDeltaStates &input_delta_states,
 
     dim3 grid_dim(grid_col, grid_row);
     dim3 block_dim(threads, threads);
+
     if (state_udapte) {
-        linear_bwd_delta_z<<<grid_dim, block_dim>>>(
-            this->d_mu_w, cu_next_bwd_states->d_jcb,
-            cu_input_delta_states->d_delta_mu,
-            cu_input_delta_states->d_delta_var, this->input_size,
-            this->output_size, batch_size, cu_output_delta_states->d_delta_mu,
-            cu_output_delta_states->d_delta_var);
+        linear_state_backward_cuda(
+            cu_input_delta_states, cu_output_delta_states, cu_next_bwd_states,
+            this->d_mu_w, this->input_size, this->output_size, batch_size);
     }
 
     // Updated values for weights
     if (this->param_update) {
-        unsigned int grid_row_w = (this->input_size + threads - 1) / threads;
-        unsigned int grid_col_w = (this->output_size + threads - 1) / threads;
-        dim3 grid_dim_w(grid_col_w, grid_row_w);
-
-        linear_bwd_delta_w<<<grid_dim_w, block_dim>>>(
-            this->d_var_w, cu_next_bwd_states->d_mu_a,
-            cu_input_delta_states->d_delta_mu,
-            cu_input_delta_states->d_delta_var, this->input_size,
-            this->output_size, batch_size, this->d_delta_mu_w,
-            this->d_delta_var_w);
+        linear_weight_backward_cuda(
+            cu_input_delta_states, cu_output_delta_states, cu_next_bwd_states,
+            this->d_var_w, this->input_size, this->output_size, batch_size,
+            this->d_delta_mu_w, this->d_delta_var_w);
 
         // Updated values for biases
         if (this->bias) {
diff --git a/src/lstm_layer.cpp b/src/lstm_layer.cpp
index ddcd4794d..98e2fde5f 100644
--- a/src/lstm_layer.cpp
+++ b/src/lstm_layer.cpp
@@ -1274,105 +1274,6 @@ void LSTM::forward(BaseHiddenStates &input_states,
     }
 }
 
-void LSTM::state_backward(BaseBackwardStates &next_bwd_states,
-                          BaseDeltaStates &input_delta_states,
-                          BaseDeltaStates &output_hidden_states,
-                          BaseTempStates &temp_states)
-/*
- */
-{
-    int batch_size = input_delta_states.block_size;
-
-    if (this->num_threads > 1) {
-        lstm_delta_mean_var_z_mp(
-            this->mu_w, lstm_states.jcb_f_ga, lstm_states.mu_i_ga,
-            lstm_states.jcb_i_ga, lstm_states.mu_c_ga, lstm_states.jcb_c_ga,
-            lstm_states.mu_o_ga, lstm_states.jcb_o_ga, lstm_states.mu_c_prev,
-            lstm_states.mu_ca, lstm_states.jcb_ca, input_delta_states.delta_mu,
-            input_delta_states.delta_var, this->w_pos_f, this->w_pos_i,
-            this->w_pos_c, this->w_pos_o, this->output_size, this->input_size,
-            this->seq_len, batch_size, this->num_threads,
-            output_hidden_states.delta_mu, output_hidden_states.delta_var);
-    } else {
-        int end_chunk = batch_size * this->seq_len * this->input_size;
-        lstm_delta_mean_var_z_worker(
-            this->mu_w, lstm_states.jcb_f_ga, lstm_states.mu_i_ga,
-            lstm_states.jcb_i_ga, lstm_states.mu_c_ga, lstm_states.jcb_c_ga,
-            lstm_states.mu_o_ga, lstm_states.jcb_o_ga, lstm_states.mu_c_prev,
-            lstm_states.mu_ca, lstm_states.jcb_ca, input_delta_states.delta_mu,
-            input_delta_states.delta_var, this->w_pos_f, this->w_pos_i,
-            this->w_pos_c, this->w_pos_o, this->output_size, this->input_size,
-            this->seq_len, 0, end_chunk, output_hidden_states.delta_mu,
-            output_hidden_states.delta_var);
-    }
-}
-
-void LSTM::param_backward(BaseBackwardStates &next_bwd_states,
-                          BaseDeltaStates &delta_states,
-                          BaseTempStates &temp_states)
-/*
- */
-{
-    int batch_size = delta_states.block_size;
-
-    if (this->num_threads > 1) {
-        lstm_cat_activations_and_prev_states_mp(
-            next_bwd_states.mu_a, lstm_states.mu_h_prev, this->input_size,
-            this->output_size, this->seq_len, batch_size, this->num_threads,
-            lstm_states.mu_ha);
-
-        lstm_delta_mean_var_w_mp(
-            this->var_w, lstm_states.mu_ha, lstm_states.jcb_f_ga,
-            lstm_states.mu_i_ga, lstm_states.jcb_i_ga, lstm_states.mu_c_ga,
-            lstm_states.jcb_c_ga, lstm_states.mu_o_ga, lstm_states.jcb_o_ga,
-            lstm_states.mu_c_prev, lstm_states.mu_ca, lstm_states.jcb_ca,
-            delta_states.delta_mu, delta_states.delta_var, this->w_pos_f,
-            this->w_pos_i, this->w_pos_c, this->w_pos_o, this->output_size,
-            this->input_size, this->seq_len, batch_size, this->num_threads,
-            this->delta_mu_w, this->delta_var_w);
-
-        if (this->bias) {
-            lstm_delta_mean_var_b_mp(
-                this->var_b, lstm_states.jcb_f_ga, lstm_states.mu_i_ga,
-                lstm_states.jcb_i_ga, lstm_states.mu_c_ga, lstm_states.jcb_c_ga,
-                lstm_states.mu_o_ga, lstm_states.jcb_o_ga,
-                lstm_states.mu_c_prev, lstm_states.mu_ca, lstm_states.jcb_ca,
-                delta_states.delta_mu, delta_states.delta_var, this->b_pos_f,
-                this->b_pos_i, this->b_pos_c, this->b_pos_o, this->output_size,
-                this->seq_len, batch_size, this->num_threads, this->delta_mu_b,
-                this->delta_var_b);
-        }
-    } else {
-        int end_chunk_w =
-            (this->input_size + this->output_size) * this->output_size;
-        lstm_cat_activations_and_prev_states(
-            next_bwd_states.mu_a, lstm_states.mu_h_prev, this->input_size,
-            this->output_size, this->seq_len, batch_size, lstm_states.mu_ha);
-
-        lstm_delta_mean_var_w_worker(
-            this->var_w, lstm_states.mu_ha, lstm_states.jcb_f_ga,
-            lstm_states.mu_i_ga, lstm_states.jcb_i_ga, lstm_states.mu_c_ga,
-            lstm_states.jcb_c_ga, lstm_states.mu_o_ga, lstm_states.jcb_o_ga,
-            lstm_states.mu_c_prev, lstm_states.mu_ca, lstm_states.jcb_ca,
-            delta_states.delta_mu, delta_states.delta_var, this->w_pos_f,
-            this->w_pos_i, this->w_pos_c, this->w_pos_o, this->output_size,
-            this->input_size, this->seq_len, batch_size, 0, end_chunk_w,
-            this->delta_mu_w, this->delta_var_w);
-
-        if (this->bias) {
-            lstm_delta_mean_var_b_worker(
-                this->var_b, lstm_states.jcb_f_ga, lstm_states.mu_i_ga,
-                lstm_states.jcb_i_ga, lstm_states.mu_c_ga, lstm_states.jcb_c_ga,
-                lstm_states.mu_o_ga, lstm_states.jcb_o_ga,
-                lstm_states.mu_c_prev, lstm_states.mu_ca, lstm_states.jcb_ca,
-                delta_states.delta_mu, delta_states.delta_var, this->b_pos_f,
-                this->b_pos_i, this->b_pos_c, this->b_pos_o, this->output_size,
-                this->seq_len, batch_size, 0, this->output_size,
-                this->delta_mu_b, this->delta_var_b);
-        }
-    }
-}
-
 void LSTM::backward(BaseDeltaStates &input_delta_states,
                     BaseDeltaStates &output_delta_states,
                     BaseTempStates &temp_states, bool state_udapte)
diff --git a/src/lstm_layer_cuda.cu b/src/lstm_layer_cuda.cu
index d5e302ea0..7c0608f59 100644
--- a/src/lstm_layer_cuda.cu
+++ b/src/lstm_layer_cuda.cu
@@ -871,101 +871,6 @@ void LSTMCuda::forward(BaseHiddenStates &input_states,
     }
 }
 
-void LSTMCuda::state_backward(BaseBackwardStates &next_bwd_states,
-                              BaseDeltaStates &input_delta_states,
-                              BaseDeltaStates &output_delta_states,
-                              BaseTempStates &temp_states)
-/*
- */
-{
-    // New poitner will point to the same memory location when casting
-    DeltaStateCuda *cu_input_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&input_delta_states);
-    DeltaStateCuda *cu_output_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&output_delta_states);
-
-    // Initialization
-    int batch_size = input_delta_states.block_size;
-
-    unsigned int gridRow_cov =
-        (batch_size * this->seq_len + this->num_cuda_threads - 1) /
-        this->num_cuda_threads;
-    unsigned int gridCol_cov = (this->input_size + this->num_cuda_threads - 1) /
-                               this->num_cuda_threads;
-    dim3 dim_grid(gridCol_cov, gridRow_cov);
-    dim3 dim_block(this->num_cuda_threads, this->num_cuda_threads);
-
-    lstm_delta_mean_var_z<<<dim_grid, dim_block>>>(
-        this->d_mu_w, this->lstm_state.d_jcb_f_ga, this->lstm_state.d_mu_i_ga,
-        this->lstm_state.d_jcb_i_ga, this->lstm_state.d_mu_c_ga,
-        this->lstm_state.d_jcb_c_ga, this->lstm_state.d_mu_o_ga,
-        this->lstm_state.d_jcb_o_ga, this->lstm_state.d_mu_c_prev,
-        this->lstm_state.d_mu_ca, this->lstm_state.d_jcb_ca,
-        cu_input_delta_states->d_delta_mu, cu_input_delta_states->d_delta_var,
-        this->w_pos_f, this->w_pos_i, this->w_pos_c, this->w_pos_o,
-        this->output_size, this->input_size, this->seq_len, batch_size,
-        cu_output_delta_states->d_delta_mu,
-        cu_output_delta_states->d_delta_var);
-}
-
-void LSTMCuda::param_backward(BaseBackwardStates &next_bwd_states,
-                              BaseDeltaStates &delta_states,
-                              BaseTempStates &temp_states)
-/*
- */
-{
-    // New poitner will point to the same memory location when casting
-    BackwardStateCuda *cu_next_bwd_states =
-        dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
-    DeltaStateCuda *cu_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&delta_states);
-
-    int batch_size = delta_states.block_size;
-    int threads = this->num_cuda_threads;
-
-    // Launch kernel
-    unsigned int b_blocks = (this->output_size + threads - 1) / threads;
-    unsigned int grid_row =
-        (this->input_size + this->output_size + threads - 1) / threads;
-    unsigned int grid_col = (this->output_size + threads - 1) / threads;
-    dim3 dim_grid(grid_col, grid_row);
-    dim3 dim_block(threads, threads);
-
-    // Concatenate the hidden states from the previous time step and
-    // activations from the previous layer
-    unsigned int grid_row_cat = (batch_size + threads - 1) / threads;
-    unsigned int grid_col_cat = (this->seq_len + threads - 1) / threads;
-    dim3 dim_grid_cat(grid_col_cat, grid_row_cat);
-    lstm_cat_act_and_prev_states_cuda<<<dim_grid_cat, dim_block>>>(
-        cu_next_bwd_states->d_mu_a, this->lstm_state.d_mu_h_prev,
-        this->input_size, this->output_size, this->seq_len, batch_size,
-        this->lstm_state.d_mu_ha);
-
-    lstm_delta_mean_var_w<<<dim_grid, dim_block>>>(
-        this->d_var_w, this->lstm_state.d_mu_ha, this->lstm_state.d_jcb_f_ga,
-        this->lstm_state.d_mu_i_ga, this->lstm_state.d_jcb_i_ga,
-        this->lstm_state.d_mu_c_ga, this->lstm_state.d_jcb_c_ga,
-        this->lstm_state.d_mu_o_ga, this->lstm_state.d_jcb_o_ga,
-        this->lstm_state.d_mu_c_prev, this->lstm_state.d_mu_ca,
-        this->lstm_state.d_jcb_ca, cu_delta_states->d_delta_mu,
-        cu_delta_states->d_delta_var, this->w_pos_f, this->w_pos_i,
-        this->w_pos_c, this->w_pos_o, this->output_size, this->input_size,
-        this->seq_len, batch_size, this->d_delta_mu_w, this->d_delta_var_w);
-
-    if (this->bias) {
-        lstm_delta_mean_var_b<<<b_blocks, threads>>>(
-            this->d_var_b, this->lstm_state.d_jcb_f_ga,
-            this->lstm_state.d_mu_i_ga, this->lstm_state.d_jcb_i_ga,
-            this->lstm_state.d_mu_c_ga, this->lstm_state.d_jcb_c_ga,
-            this->lstm_state.d_mu_o_ga, this->lstm_state.d_jcb_o_ga,
-            this->lstm_state.d_mu_c_prev, this->lstm_state.d_mu_ca,
-            this->lstm_state.d_jcb_ca, cu_delta_states->d_delta_mu,
-            cu_delta_states->d_delta_var, this->b_pos_f, this->b_pos_i,
-            this->b_pos_c, this->b_pos_o, this->output_size, this->seq_len,
-            batch_size, this->d_delta_mu_b, this->d_delta_var_b);
-    }
-}
-
 void LSTMCuda::backward(BaseDeltaStates &input_delta_states,
                         BaseDeltaStates &output_delta_states,
                         BaseTempStates &temp_states, bool state_udapte)
diff --git a/src/norm_layer.cpp b/src/norm_layer.cpp
index 39ee0d40c..942e2677e 100644
--- a/src/norm_layer.cpp
+++ b/src/norm_layer.cpp
@@ -1570,136 +1570,6 @@ void LayerNorm::forward(BaseHiddenStates &input_states,
     }
 }
 
-void LayerNorm::state_backward(BaseBackwardStates &next_bwd_states,
-                               BaseDeltaStates &input_delta_states,
-                               BaseDeltaStates &output_delta_states,
-                               BaseTempStates &temp_states)
-/*
- */
-{
-    int batch_size = input_delta_states.block_size;
-    if (this->num_threads <= 1) {
-        if (this->normalized_shape.size() == 1) {
-            layernorm_bwd_delta_z(
-                this->mu_w, next_bwd_states.jcb, this->var_ra,
-                input_delta_states.delta_mu, input_delta_states.delta_var,
-                this->epsilon, this->input_size, 0, batch_size,
-                output_delta_states.delta_mu, output_delta_states.delta_var);
-        } else {
-            int wihi = this->in_height * this->in_width;
-
-            layernorm2d_bwd_delta_z(
-                this->mu_w, next_bwd_states.jcb, this->var_ra,
-                input_delta_states.delta_mu, input_delta_states.delta_var,
-                this->epsilon, wihi, this->in_channels, 0, batch_size,
-                output_delta_states.delta_mu, output_delta_states.delta_var);
-        }
-    } else {
-        if (this->normalized_shape.size() == 1) {
-            layernorm_bwd_delta_z_mp(
-                this->mu_w, next_bwd_states.jcb, this->var_ra,
-                input_delta_states.delta_mu, input_delta_states.delta_var,
-                this->epsilon, this->input_size, batch_size, this->num_threads,
-                output_delta_states.delta_mu, output_delta_states.delta_var);
-        } else {
-            int wihi = this->in_height * this->in_width;
-
-            layernorm2d_bwd_delta_z_mp(
-                this->mu_w, next_bwd_states.jcb, this->var_ra,
-                input_delta_states.delta_mu, input_delta_states.delta_var,
-                this->epsilon, wihi, this->in_channels, batch_size,
-                this->num_threads, output_delta_states.delta_mu,
-                output_delta_states.delta_var);
-        }
-    }
-}
-
-void LayerNorm::param_backward(BaseBackwardStates &next_bwd_states,
-                               BaseDeltaStates &delta_states,
-                               BaseTempStates &temp_states)
-/*
- */
-{
-    int batch_size = delta_states.block_size;
-    if (this->num_threads <= 1) {
-        if (this->normalized_shape.size() == 1) {
-            layernorm_bwd_delta_w(
-                this->var_w, next_bwd_states.mu_a, this->mu_ra, this->var_ra,
-                delta_states.delta_mu, delta_states.delta_var, this->epsilon,
-                this->input_size, batch_size, 0, this->input_size,
-                this->delta_mu_w, this->delta_var_w);
-
-            if (this->bias) {
-                layernorm_bwd_delta_b(
-                    this->var_b, delta_states.delta_mu, delta_states.delta_var,
-                    this->epsilon, this->input_size, batch_size, 0,
-                    this->input_size, this->delta_mu_b, this->delta_var_b);
-            }
-        } else {
-            int wihi = this->in_height * this->in_width;
-
-            layernorm2d_bwd_delta_w(
-                this->var_w, next_bwd_states.mu_a, this->mu_ra, this->var_ra,
-                delta_states.delta_mu, delta_states.delta_var, this->epsilon,
-                wihi, this->in_channels, 0, batch_size, temp_states.tmp_1,
-                temp_states.tmp_2);
-
-            delta_param_sum(temp_states.tmp_1, temp_states.tmp_2, wihi,
-                            this->in_channels, batch_size, this->delta_mu_w,
-                            this->delta_var_w);
-
-            if (this->bias) {
-                layernorm2d_bwd_delta_b(this->var_b, delta_states.delta_mu,
-                                        delta_states.delta_var, this->epsilon,
-                                        wihi, this->in_channels, 0, batch_size,
-                                        temp_states.tmp_1, temp_states.tmp_2);
-
-                delta_param_sum(temp_states.tmp_1, temp_states.tmp_2, wihi,
-                                this->in_channels, batch_size, this->delta_mu_b,
-                                this->delta_var_b);
-            }
-        }
-    } else {
-        if (this->normalized_shape.size() == 1) {
-            layernorm_bwd_delta_w_mp(
-                this->var_w, next_bwd_states.mu_a, this->mu_ra, this->var_ra,
-                delta_states.delta_mu, delta_states.delta_var, this->epsilon,
-                this->input_size, batch_size, this->num_threads,
-                this->delta_mu_w, this->delta_var_w);
-
-            if (this->bias) {
-                layernorm_bwd_delta_b_mp(
-                    this->var_b, delta_states.delta_mu, delta_states.delta_var,
-                    this->epsilon, this->input_size, batch_size,
-                    this->num_threads, this->delta_mu_b, this->delta_var_b);
-            }
-        } else {
-            int wihi = this->in_height * this->in_width;
-
-            layernorm2d_bwd_delta_w_mp(
-                this->var_w, next_bwd_states.mu_a, this->mu_ra, this->var_ra,
-                delta_states.delta_mu, delta_states.delta_var, this->epsilon,
-                wihi, this->in_channels, batch_size, this->num_threads,
-                temp_states.tmp_1, temp_states.tmp_2);
-
-            delta_param_sum(temp_states.tmp_1, temp_states.tmp_2, wihi,
-                            this->in_channels, batch_size, this->delta_mu_w,
-                            this->delta_var_w);
-
-            if (this->bias) {
-                layernorm2d_bwd_delta_b_mp(
-                    this->var_b, delta_states.delta_mu, delta_states.delta_var,
-                    this->epsilon, wihi, this->in_channels, batch_size,
-                    this->num_threads, temp_states.tmp_1, temp_states.tmp_2);
-
-                delta_param_sum(temp_states.tmp_1, temp_states.tmp_2, wihi,
-                                this->in_channels, batch_size, this->delta_mu_b,
-                                this->delta_var_b);
-            }
-        }
-    }
-}
-
 void LayerNorm::backward(BaseDeltaStates &input_delta_states,
                          BaseDeltaStates &output_delta_states,
                          BaseTempStates &temp_states, bool state_udapte)
@@ -2159,148 +2029,6 @@ void BatchNorm2d::forward(BaseHiddenStates &input_states,
     }
 }
 
-void BatchNorm2d::state_backward(BaseBackwardStates &next_bwd_states,
-                                 BaseDeltaStates &input_delta_states,
-                                 BaseDeltaStates &output_delta_states,
-                                 BaseTempStates &temp_states)
-/*
- */
-{
-    int batch_size = input_delta_states.block_size;
-
-    if (this->num_threads == 1) {
-        if (this->in_channels == 0) {
-            batchnorm_bwd_delta_z(
-                this->mu_w, next_bwd_states.jcb, this->var_norm_batch,
-                input_delta_states.delta_mu, input_delta_states.delta_var,
-                this->epsilon, this->input_size, 0, batch_size,
-                output_delta_states.delta_mu, output_delta_states.delta_var);
-        } else {
-            int wihi = this->in_width * this->in_height;
-            int end_chunk = this->in_channels * batch_size;
-
-            batchnorm2d_bwd_delta_z(
-                this->mu_w, next_bwd_states.jcb, this->var_norm_batch,
-                input_delta_states.delta_mu, input_delta_states.delta_var,
-                this->epsilon, wihi, this->in_channels, batch_size, 0,
-                end_chunk, output_delta_states.delta_mu,
-                output_delta_states.delta_var);
-        }
-    } else {
-        if (this->in_channels == 0) {
-            batchnorm_bwd_delta_z_mp(
-                this->mu_w, next_bwd_states.jcb, this->var_norm_batch,
-                input_delta_states.delta_mu, input_delta_states.delta_var,
-                this->epsilon, this->input_size, batch_size, this->num_threads,
-                output_delta_states.delta_mu, output_delta_states.delta_var);
-        } else {
-            int wihi = this->in_width * this->in_height;
-            int end_chunk = this->in_channels * batch_size;
-
-            batchnorm2d_bwd_delta_z_mp(
-                this->mu_w, next_bwd_states.jcb, this->var_norm_batch,
-                input_delta_states.delta_mu, input_delta_states.delta_var,
-                this->epsilon, wihi, this->in_channels, batch_size,
-                this->num_threads, output_delta_states.delta_mu,
-                output_delta_states.delta_var);
-        }
-    }
-}
-
-void BatchNorm2d::param_backward(BaseBackwardStates &next_bwd_states,
-                                 BaseDeltaStates &delta_states,
-                                 BaseTempStates &temp_states)
-/*
- */
-{
-    int batch_size = delta_states.block_size;
-
-    if (this->num_threads == 1) {
-        if (this->in_channels == 0) {
-            batchnorm_bwd_delta_w(this->var_w, next_bwd_states.mu_a,
-                                  this->mu_norm_batch, this->var_norm_batch,
-                                  delta_states.delta_mu, delta_states.delta_var,
-                                  this->epsilon, this->input_size, batch_size,
-                                  0, this->input_size, this->delta_mu_w,
-                                  this->delta_var_w);
-
-            if (this->bias) {
-                batchnorm_bwd_delta_b(
-                    this->var_b, delta_states.delta_mu, delta_states.delta_var,
-                    this->epsilon, this->input_size, batch_size, 0,
-                    this->input_size, this->delta_mu_b, this->delta_var_b);
-            }
-
-        } else {
-            int wihi = this->in_width * this->in_height;
-            int end_chunk = this->in_channels * batch_size;
-
-            batchnorm2d_bwd_delta_w(
-                this->var_w, next_bwd_states.mu_a, this->mu_norm_batch,
-                this->var_norm_batch, delta_states.delta_mu,
-                delta_states.delta_var, this->epsilon, wihi, this->in_channels,
-                batch_size, 0, end_chunk, temp_states.tmp_1, temp_states.tmp_2);
-
-            delta_param_sum(temp_states.tmp_1, temp_states.tmp_2, wihi,
-                            this->in_channels, batch_size, this->delta_mu_w,
-                            this->delta_var_w);
-
-            if (this->num_biases > 0) {
-                batchnorm2d_bwd_delta_b(this->var_b, delta_states.delta_mu,
-                                        delta_states.delta_var, this->epsilon,
-                                        wihi, this->in_channels, 0, end_chunk,
-                                        temp_states.tmp_1, temp_states.tmp_2);
-
-                delta_param_sum(temp_states.tmp_1, temp_states.tmp_2, wihi,
-                                this->in_channels, batch_size, this->delta_mu_b,
-                                this->delta_var_b);
-            }
-        }
-    } else {
-        if (this->in_channels == 0) {
-            batchnorm_bwd_delta_w_mp(
-                this->var_w, next_bwd_states.mu_a, this->mu_norm_batch,
-                this->var_norm_batch, delta_states.delta_mu,
-                delta_states.delta_var, this->epsilon, this->input_size,
-                batch_size, this->num_threads, this->delta_mu_w,
-                this->delta_var_w);
-
-            if (this->bias) {
-                batchnorm_bwd_delta_b_mp(
-                    this->var_b, delta_states.delta_mu, delta_states.delta_var,
-                    this->epsilon, this->input_size, batch_size,
-                    this->num_threads, this->delta_mu_b, this->delta_var_b);
-            }
-
-        } else {
-            int wihi = this->in_width * this->in_height;
-            int end_chunk = this->in_channels * batch_size;
-
-            batchnorm2d_bwd_delta_w_mp(
-                this->var_w, next_bwd_states.mu_a, this->mu_norm_batch,
-                this->var_norm_batch, delta_states.delta_mu,
-                delta_states.delta_var, this->epsilon, wihi, this->in_channels,
-                batch_size, this->num_threads, temp_states.tmp_1,
-                temp_states.tmp_2);
-
-            delta_param_sum(temp_states.tmp_1, temp_states.tmp_2, wihi,
-                            this->in_channels, batch_size, this->delta_mu_w,
-                            this->delta_var_w);
-
-            if (this->num_biases > 0) {
-                batchnorm2d_bwd_delta_b_mp(
-                    this->var_b, delta_states.delta_mu, delta_states.delta_var,
-                    this->epsilon, wihi, this->in_channels, batch_size,
-                    this->num_threads, temp_states.tmp_1, temp_states.tmp_2);
-
-                delta_param_sum(temp_states.tmp_1, temp_states.tmp_2, wihi,
-                                this->in_channels, batch_size, this->delta_mu_b,
-                                this->delta_var_b);
-            }
-        }
-    }
-}
-
 void BatchNorm2d::backward(BaseDeltaStates &input_delta_states,
                            BaseDeltaStates &output_delta_states,
                            BaseTempStates &temp_states, bool state_udapte)
diff --git a/src/norm_layer_cuda.cu b/src/norm_layer_cuda.cu
index dba9064d5..92f613a96 100644
--- a/src/norm_layer_cuda.cu
+++ b/src/norm_layer_cuda.cu
@@ -1143,120 +1143,6 @@ void LayerNormCuda::forward(BaseHiddenStates &input_states,
     }
 }
 
-void LayerNormCuda::state_backward(BaseBackwardStates &next_bwd_states,
-                                   BaseDeltaStates &input_delta_states,
-                                   BaseDeltaStates &output_delta_states,
-                                   BaseTempStates &temp_states)
-/*
- */
-{
-    // New poitner will point to the same memory location when casting
-    BackwardStateCuda *cu_next_bwd_states =
-        dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
-    DeltaStateCuda *cu_input_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&input_delta_states);
-    DeltaStateCuda *cu_output_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&output_delta_states);
-
-    // Initialization
-    int batch_size = input_delta_states.block_size;
-    int num_threads = this->num_cuda_threads;
-    dim3 block_dim(num_threads, num_threads);
-
-    unsigned int grid_row = (batch_size + num_threads - 1) / num_threads;
-    unsigned int grid_col = (this->input_size + num_threads - 1) / num_threads;
-    dim3 grid_size(grid_col, grid_row);
-
-    if (this->normalized_shape.size() == 1) {
-        layernorm_bwd_delta_z_cuda<<<grid_size, block_dim>>>(
-            this->d_mu_w, cu_next_bwd_states->d_jcb, this->d_var_ra,
-            cu_input_delta_states->d_delta_mu,
-            cu_input_delta_states->d_delta_var, this->epsilon, this->input_size,
-            batch_size, cu_output_delta_states->d_delta_mu,
-            cu_output_delta_states->d_delta_var);
-    } else {
-        int wihi = this->in_height * this->in_width;
-
-        layernorm2d_bwd_delta_z_cuda<<<grid_size, block_dim>>>(
-            this->d_mu_w, cu_next_bwd_states->d_jcb, this->d_var_ra,
-            cu_input_delta_states->d_delta_mu,
-            cu_input_delta_states->d_delta_var, this->epsilon, wihi,
-            this->in_channels, batch_size, cu_output_delta_states->d_delta_mu,
-            cu_output_delta_states->d_delta_var);
-    }
-}
-
-void LayerNormCuda::param_backward(BaseBackwardStates &next_bwd_states,
-                                   BaseDeltaStates &delta_states,
-                                   BaseTempStates &temp_states)
-/*
- */
-{
-    // New poitner will point to the same memory location when casting
-    BackwardStateCuda *cu_next_bwd_states =
-        dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
-    DeltaStateCuda *cu_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&delta_states);
-    TempStateCuda *cu_temp_states = dynamic_cast<TempStateCuda *>(&temp_states);
-
-    // Initalization
-    int batch_size = delta_states.block_size;
-    int num_threads = this->num_cuda_threads;
-    dim3 block_dim(num_threads, num_threads);
-
-    unsigned int grid_col = (this->input_size + num_threads - 1) / num_threads;
-
-    if (this->normalized_shape.size() == 1) {
-        layernorm_bwd_delta_w_cuda<<<grid_col, num_threads>>>(
-            this->d_var_w, cu_next_bwd_states->d_mu_a, this->d_mu_ra,
-            this->d_var_ra, cu_delta_states->d_delta_mu,
-            cu_delta_states->d_delta_var, this->epsilon, this->input_size,
-            batch_size, this->d_delta_mu_w, this->d_delta_var_w);
-
-        if (this->bias) {
-            layernorm_bwd_delta_b_cuda<<<grid_col, num_threads>>>(
-                this->d_var_b, cu_delta_states->d_delta_mu,
-                cu_delta_states->d_delta_var, this->epsilon, this->input_size,
-                batch_size, this->d_delta_mu_b, this->d_delta_var_b);
-        }
-
-    } else {
-        int wihi = this->in_height * this->in_width;
-        unsigned int grid_row = (batch_size + num_threads - 1) / num_threads;
-        dim3 grid_size(grid_col, grid_row);
-        unsigned int sum_grid_size =
-            (this->in_channels + num_threads - 1) / num_threads;
-
-        // Weights
-        // TODO: Not sure if it should be batch_size or batch_size * fi
-        layernorm2d_bwd_delta_w_cuda<<<grid_size, block_dim>>>(
-            this->d_var_w, cu_next_bwd_states->d_mu_a, this->d_mu_ra,
-            this->d_var_ra, cu_delta_states->d_delta_mu,
-            cu_delta_states->d_delta_var, this->epsilon, wihi,
-            this->in_channels, batch_size, cu_temp_states->d_tmp_1,
-            cu_temp_states->d_tmp_2);
-
-        delta_param_sum<<<sum_grid_size, num_threads>>>(
-            cu_temp_states->d_tmp_1, cu_temp_states->d_tmp_2, wihi,
-            this->in_channels, batch_size, this->d_delta_mu_w,
-            this->d_delta_var_w);
-
-        // Biases
-        if (this->bias) {
-            layernorm2d_bwd_delta_b_cuda<<<grid_size, block_dim>>>(
-                this->d_var_b, cu_delta_states->d_delta_mu,
-                cu_delta_states->d_delta_var, this->epsilon, wihi,
-                this->in_channels, batch_size, cu_temp_states->d_tmp_1,
-                cu_temp_states->d_tmp_2);
-
-            delta_param_sum<<<sum_grid_size, num_threads>>>(
-                cu_temp_states->d_tmp_1, cu_temp_states->d_tmp_2, wihi,
-                this->in_channels, batch_size, this->d_delta_mu_b,
-                this->d_delta_var_b);
-        }
-    }
-}
-
 void LayerNormCuda::backward(BaseDeltaStates &input_delta_states,
                              BaseDeltaStates &output_delta_states,
                              BaseTempStates &temp_states, bool state_udapte)
@@ -1536,6 +1422,17 @@ void BatchNorm2dCuda::init_weight_bias()
     this->params_to_device();
 }
 
+void BatchNorm2dCuda::deallocate_running_mean_var() {
+    cudaFree(this->d_mu_ra);
+    cudaFree(this->d_var_ra);
+    cudaFree(this->d_mu_norm_batch);
+    cudaFree(this->d_var_norm_batch);
+    this->d_mu_ra = nullptr;
+    this->d_var_ra = nullptr;
+    this->d_mu_norm_batch = nullptr;
+    this->d_var_norm_batch = nullptr;
+}
+
 void BatchNorm2dCuda::allocate_running_mean_var()
 /*
  */
@@ -1723,21 +1620,6 @@ void BatchNorm2dCuda::forward(BaseHiddenStates &input_states,
     }
 }
 
-void BatchNorm2dCuda::state_backward(BaseBackwardStates &next_bwd_states,
-                                     BaseDeltaStates &input_delta_states,
-                                     BaseDeltaStates &output_delta_states,
-                                     BaseTempStates &temp_states)
-/*
- */
-{}
-
-void BatchNorm2dCuda::param_backward(BaseBackwardStates &next_bwd_states,
-                                     BaseDeltaStates &delta_states,
-                                     BaseTempStates &temp_states)
-/*
- */
-{}
-
 void BatchNorm2dCuda::backward(BaseDeltaStates &input_delta_states,
                                BaseDeltaStates &output_delta_states,
                                BaseTempStates &temp_states, bool state_udapte)
diff --git a/src/pooling_layer.cpp b/src/pooling_layer.cpp
index 3ef42dbc9..0f19804a5 100644
--- a/src/pooling_layer.cpp
+++ b/src/pooling_layer.cpp
@@ -120,74 +120,6 @@ void AvgPool2d::forward(BaseHiddenStates &input_states,
     }
 }
 
-void AvgPool2d::state_backward(BaseBackwardStates &next_bwd_states,
-                               BaseDeltaStates &input_delta_states,
-                               BaseDeltaStates &output_delta_states,
-                               BaseTempStates &temp_states)
-/*
- */
-{
-    // Initialization
-    int batch_size = input_delta_states.block_size;
-
-    // Launch kernel
-    int woho = this->out_width * this->out_height;
-    int wihi = this->in_width * this->in_height;
-    int pad_out_idx = woho * this->out_channels * batch_size + 1;
-
-    if (this->num_threads > 1) {
-        if (this->overlap) {
-            int num_in_states = this->in_width * this->in_height *
-                                this->in_channels * batch_size;
-
-            avgpool2d_bwd_overlapped_delta_z_mp(
-                next_bwd_states.jcb, input_delta_states.delta_mu,
-                input_delta_states.delta_var, this->z_ud_idx, woho, wihi,
-                this->kernel_size, this->col_z_ud, num_in_states, pad_out_idx,
-                this->num_threads, output_delta_states.delta_mu,
-                output_delta_states.delta_var);
-        } else {
-            int kiwo = this->kernel_size * this->out_width;
-            int nums = wihi * this->in_channels * batch_size / kiwo;
-
-            avgpool2d_bwd_delta_z_mp(
-                next_bwd_states.jcb, input_delta_states.delta_mu,
-                input_delta_states.delta_var, this->out_channels,
-                this->kernel_size, nums, this->num_threads,
-                output_delta_states.delta_mu, output_delta_states.delta_var);
-        }
-    } else {
-        if (this->overlap) {
-            int num_in_states = this->in_width * this->in_height *
-                                this->in_channels * batch_size;
-
-            avgpool2d_bwd_overlapped_delta_z(
-                next_bwd_states.jcb, input_delta_states.delta_mu,
-                input_delta_states.delta_var, this->z_ud_idx, woho, wihi,
-                this->kernel_size, this->col_z_ud, num_in_states, pad_out_idx,
-                0, num_in_states, output_delta_states.delta_mu,
-                output_delta_states.delta_var);
-        } else {
-            int kiwo = this->kernel_size * this->out_width;
-            int nums = wihi * this->in_channels * batch_size / kiwo;
-            int end_chunk = this->kernel_size * this->out_width * nums;
-
-            avgpool2d_bwd_delta_z(
-                next_bwd_states.jcb, input_delta_states.delta_mu,
-                input_delta_states.delta_var, this->out_channels,
-                this->kernel_size, nums, 0, end_chunk,
-                output_delta_states.delta_mu, output_delta_states.delta_var);
-        }
-    }
-}
-
-void AvgPool2d::param_backward(BaseBackwardStates &next_bwd_states,
-                               BaseDeltaStates &delta_states,
-                               BaseTempStates &temp_states)
-/*
- */
-{}
-
 void AvgPool2d::backward(BaseDeltaStates &input_delta_states,
                          BaseDeltaStates &output_delta_states,
                          BaseTempStates &temp_states, bool state_udapte)
diff --git a/src/pooling_layer_cuda.cu b/src/pooling_layer_cuda.cu
index ea4209290..a72da2923 100644
--- a/src/pooling_layer_cuda.cu
+++ b/src/pooling_layer_cuda.cu
@@ -113,65 +113,6 @@ void AvgPool2dCuda::forward(BaseHiddenStates &input_states,
     }
 }
 
-void AvgPool2dCuda::state_backward(BaseBackwardStates &next_bwd_states,
-                                   BaseDeltaStates &input_delta_states,
-                                   BaseDeltaStates &output_delta_states,
-                                   BaseTempStates &temp_states)
-/*
- */
-{
-    // New poitner will point to the same memory location when casting
-    BackwardStateCuda *cu_next_bwd_states =
-        dynamic_cast<BackwardStateCuda *>(&next_bwd_states);
-    DeltaStateCuda *cu_input_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&input_delta_states);
-    DeltaStateCuda *cu_output_delta_states =
-        dynamic_cast<DeltaStateCuda *>(&output_delta_states);
-
-    // Initialization
-    int batch_size = input_delta_states.block_size;
-    unsigned int num_threads = this->num_cuda_threads;
-
-    // Launch kernel
-    int woho = this->out_width * this->out_height;
-    int wihi = this->in_width * this->in_height;
-    int pad_out_idx = woho * this->out_channels * batch_size + 1;
-    if (overlap) {
-        int num_in_states =
-            this->in_width * this->in_height * this->in_channels * batch_size;
-        unsigned int grid_size =
-            (num_in_states + num_threads - 1) / num_threads;
-
-        avgpool2d_bwd_overlapped_delta_z_cuda<<<grid_size, num_threads>>>(
-            cu_next_bwd_states->d_jcb, cu_input_delta_states->d_delta_mu,
-            cu_input_delta_states->d_delta_var, this->d_z_ud_idx, woho, wihi,
-            this->kernel_size, this->col_z_ud, num_in_states, pad_out_idx,
-            cu_output_delta_states->d_delta_mu,
-            cu_output_delta_states->d_delta_var);
-
-    } else {
-        int kiwo = this->kernel_size * this->out_width;
-        int nums = wihi * this->in_channels * batch_size / kiwo;
-        unsigned int grid_row = (kiwo + num_threads - 1) / num_threads;
-        unsigned int grid_col = (nums + num_threads - 1) / num_threads;
-        dim3 dim_grid(grid_col, grid_row);
-        dim3 dim_block(num_threads, num_threads);
-
-        avgpool2d_bwd_delta_z_cuda<<<dim_grid, grid_row>>>(
-            cu_next_bwd_states->d_jcb, cu_input_delta_states->d_delta_mu,
-            cu_input_delta_states->d_delta_var, this->out_width,
-            this->kernel_size, nums, cu_output_delta_states->d_delta_mu,
-            cu_output_delta_states->d_delta_var);
-    }
-}
-
-void AvgPool2dCuda::param_backward(BaseBackwardStates &next_bwd_states,
-                                   BaseDeltaStates &delta_states,
-                                   BaseTempStates &temp_states)
-/*
- */
-{}
-
 void AvgPool2dCuda::backward(BaseDeltaStates &input_delta_states,
                              BaseDeltaStates &output_delta_states,
                              BaseTempStates &temp_states, bool state_udapte)


Test patch:
diff --git a/test/fnn/test_fnn_mnist_cpu.cpp b/test/fnn/test_fnn_mnist_cpu.cpp
index a988f5849..cb6c256fa 100644
--- a/test/fnn/test_fnn_mnist_cpu.cpp
+++ b/test/fnn/test_fnn_mnist_cpu.cpp
@@ -66,12 +66,12 @@ void fnn_mnist() {
     //////////////////////////////////////////////////////////////////////
     // TAGI network
     //////////////////////////////////////////////////////////////////////
-    // Sequential model(Linear(784, 100), ReLU(), Linear(100, 100), ReLU(),
-    //                  Linear(100, 11));
+    Sequential model(Linear(784, 4096), ReLU(), Linear(4096, 4096), ReLU(),
+                     Linear(4096, 11));
 
-    // Sequential model(Linear(784, 100), BatchNorm2d(100), ReLU(),
-    //                  Linear(100, 100), BatchNorm2d(100), ReLU(),
-    //                  Linear(100, 11));
+    // Sequential model(Linear(784, 1024), BatchNorm2d(1024), ReLU(),
+    //                  Linear(1024, 1024), BatchNorm2d(1024), ReLU(),
+    //                  Linear(1024, 11));
 
     // Sequential model(Linear(784, 100), LayerNorm(std::vector<int>({100})),
     //                  ReLU(), Linear(100, 100),
@@ -89,12 +89,12 @@ void fnn_mnist() {
     //                  BatchNorm2d(32), ReLU(), AvgPool2d(3, 2),
     //                  Linear(32 * 4 * 4, 100), ReLU(), Linear(100, 11));
 
-    Sequential model(Conv2d(1, 16, 4, false, 1, 1, 1, 28, 28),
-                     LayerNorm(std::vector<int>({16, 27, 27})), ReLU(),
-                     AvgPool2d(3, 2), Conv2d(16, 32, 5, false),
-                     LayerNorm(std::vector<int>({32, 9, 9})), ReLU(),
-                     AvgPool2d(3, 2), Linear(32 * 4 * 4, 100), ReLU(),
-                     Linear(100, 11));
+    // Sequential model(Conv2d(1, 16, 4, false, 1, 1, 1, 28, 28),
+    //                  LayerNorm(std::vector<int>({16, 27, 27})), ReLU(),
+    //                  AvgPool2d(3, 2), Conv2d(16, 32, 5, false),
+    //                  LayerNorm(std::vector<int>({32, 9, 9})), ReLU(),
+    //                  AvgPool2d(3, 2), Linear(32 * 4 * 4, 100), ReLU(),
+    //                  Linear(100, 11));
 
     // model.set_threads(8);
     model.to_device("cuda");
@@ -142,8 +142,8 @@ void fnn_mnist() {
     //////////////////////////////////////////////////////////////////////
     unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
     std::default_random_engine seed_e(seed);
-    int n_epochs = 1;
-    int batch_size = 512;
+    int n_epochs = 10;
+    int batch_size = 256;
     float sigma_obs = 2.0;
     int iters = train_db.num_data / batch_size;
     std::cout << "num_iter: " << iters << "\n";
@@ -172,7 +172,7 @@ void fnn_mnist() {
         std::cout << "Epoch #" << e + 1 << "/" << n_epochs << "\n";
         std::cout << "Training...\n";
         auto start = std::chrono::steady_clock::now();
-        for (int i = 0; i < 2; i++) {
+        for (int i = 0; i < iters; i++) {
             // Load data
             get_batch_images_labels(train_db, data_idx, batch_size, i, x_batch,
                                     y_batch, idx_ud_batch, label_batch);
diff --git a/test/resnet/test_resnet_cifar10.cpp b/test/resnet/test_resnet_cifar10.cpp
index 68641fb67..7e32987c3 100644
--- a/test/resnet/test_resnet_cifar10.cpp
+++ b/test/resnet/test_resnet_cifar10.cpp
@@ -188,7 +188,7 @@ void resnet_cifar10()
     //////////////////////////////////////////////////////////////////////
     unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
     std::default_random_engine seed_e(seed);
-    int n_epochs = 1;
+    int n_epochs = 2;
     int batch_size = 128;
     float sigma_obs = 1;
     int iters = train_db.num_data / batch_size;
@@ -218,7 +218,7 @@ void resnet_cifar10()
         std::cout << "Epoch #" << e + 1 << "/" << n_epochs << "\n";
         std::cout << "Training...\n";
         auto start = std::chrono::steady_clock::now();
-        for (int i = 0; i < 200; i++) {
+        for (int i = 0; i < iters; i++) {
             // Load data
             get_batch_images_labels(train_db, data_idx, batch_size, i, x_batch,
                                     y_batch, idx_ud_batch, label_batch);