diff --git a/modeling_and_evaluation/SoC_Estimation_LSTM.ipynb b/modeling_and_evaluation/SoC_Estimation_LSTM.ipynb
index 4063ec7..e07d204 100644
--- a/modeling_and_evaluation/SoC_Estimation_LSTM.ipynb
+++ b/modeling_and_evaluation/SoC_Estimation_LSTM.ipynb
@@ -148,28 +148,21 @@
     "    def get_times(self):\n",
     "        return self.times\n",
     "\n",
-    "# LSTM Model\n",
-    "class LSTM(nn.Module):\n",
+    "# SoCLSTM Model\n",
+    "class SoCLSTM(nn.Module):\n",
     "    def __init__(self, input_size, hidden_size, num_layers):\n",
-    "        super(LSTM, self).__init__()\n",
-    "\n",
-    "        self.conv = nn.Conv1d(in_channels=1, out_channels=6, kernel_size=3, stride=1)\n",
-    "        self.relu1 = nn.ReLU()\n",
-    "        self.rnn = nn.LSTM(6, hidden_size, num_layers)\n",
-    "        self.relu2 = nn.ReLU()\n",
-    "        self.reg_1 = nn.Linear(hidden_size, 1)\n",
-    "        self.reg_2 = nn.Linear(6, 1)\n",
-    "        \n",
+    "        super(SoCLSTM, self).__init__()\n",
+    "        self.hidden_size = hidden_size\n",
+    "        self.num_layers = num_layers\n",
+    "        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)\n",
+    "        self.fc = nn.Linear(hidden_size, 1)\n",
+    "\n",
     "    def forward(self, x):\n",
-    "        x = self.relu1(self.conv(x))\n",
-    "        x, _ = self.rnn(x)\n",
-    "        s, b, h = x.shape\n",
-    "        x = x.view(s*b, h)\n",
-    "        x = self.relu2(self.reg_1(x))\n",
-    "        x = x.view(s, -1)\n",
-    "        x = self.reg_2(x)\n",
-    "        x = x.view(s, 1, 1)\n",
-    "        return x\n",
+    "        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size, dtype=x.dtype, device=x.device)\n",
+    "        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size, dtype=x.dtype, device=x.device)\n",
+    "        out, _ = self.lstm(x, (h0, c0))\n",
+    "        out = self.fc(out[:, -1, :])\n",
+    "        return out\n",
     "\n",
     "# Training loop with validation\n",
     "def train_and_validate(model, criterion, optimizer, train_loader, val_loader, epochs, device, patience=20, min_delta=0.001):\n",
@@ -3039,7 +3032,7 @@
   "kernelspec": {
    "display_name": "myenv",
    "language": "python",
-   "name": "myenv"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {