diff --git a/notebooks/dataset.ipynb b/notebooks/dataset.ipynb
index 7733be82f1d54b00aef36efb4205f9f1d0d76b3a..d8c4c5d320cbf66e7f81b02e6c1c827dbe332ab5 100644
--- a/notebooks/dataset.ipynb
+++ b/notebooks/dataset.ipynb
@@ -16,7 +16,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -25,7 +25,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -57,42 +57,42 @@
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>how are you?</td>\n",
-       "      <td>good</td>\n",
+       "      <td>sad</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
-       "      <td>how are you?</td>\n",
+       "      <td>how re you?</td>\n",
        "      <td>sad</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
-       "      <td>how are you?</td>\n",
-       "      <td>upset</td>\n",
+       "      <td>what is your name?</td>\n",
+       "      <td>john</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
-       "      <td>how old are you?</td>\n",
-       "      <td>23 years old</td>\n",
+       "      <td>are you good?</td>\n",
+       "      <td>no</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>4</th>\n",
        "      <td>how old are you?</td>\n",
-       "      <td>9 years old</td>\n",
+       "      <td>47 years old</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
-       "           question        answer\n",
-       "0      how are you?          good\n",
-       "1      how are you?           sad\n",
-       "2      how are you?         upset\n",
-       "3  how old are you?  23 years old\n",
-       "4  how old are you?   9 years old"
+       "             question        answer\n",
+       "0        how are you?           sad\n",
+       "1         how re you?           sad\n",
+       "2  what is your name?          john\n",
+       "3       are you good?            no\n",
+       "4    how old are you?  47 years old"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -118,6 +118,7 @@
    "source": [
     "import torch\n",
     "from torch.utils.data.dataset import Dataset\n",
+    "import pandas as pd\n",
     "import numpy as np\n",
     "\n",
     "\n",
@@ -127,23 +128,27 @@
     "        self.questions = self.data[\"question\"]\n",
     "        self.answers = self.data[\"answer\"]\n",
     "        self.data_len = len(self.data.index)\n",
-    "        \n",
+    "\n",
     "        # Unique characters in the database.\n",
     "        self.unique_characters = set(\"\".join(self.questions + self.answers))\n",
-    "        self.unique_characters_length = len(self.unique_characters)\n",
+    "        self.unique_characters_length = len(self.unique_characters)  # 24\n",
+    "\n",
     "        # Map int to character.\n",
-    "        self.int2char = dict(enumerate(self.unique_characters))\n",
+    "        self.int2char = {i: char for i, char in enumerate(self.unique_characters)}\n",
     "        # Map character to int.\n",
-    "        self.char2int = {char: i for i, char in self.int2char.items()}\n",
-    "        \n",
+    "        self.char2int = {char: i for i, char in enumerate(self.unique_characters)}\n",
+    "\n",
     "        # Longer question.\n",
-    "        longer_question_length = len(max(self.questions, key=len))\n",
+    "        self.longer_question_length = len(max(self.questions, key=len))  # 24\n",
     "        # Longer answer.\n",
-    "        longer_answer_length = len(max(self.answers, key=len))\n",
-    "        \n",
+    "        self.longer_answer_length = self.longer_question_length\n",
+    "        # self.longer_answer_length = len(max(self.answers, key=len))  # 14\n",
+    "\n",
     "        # Pad strings.\n",
-    "        self.questions = self.questions.str.pad(longer_question_length, side=\"right\")\n",
-    "        self.answers = self.answers.str.pad(longer_answer_length, side=\"right\")\n",
+    "        self.questions = self.questions.str.pad(\n",
+    "            self.longer_question_length, side=\"right\"\n",
+    "        )\n",
+    "        self.answers = self.answers.str.pad(self.longer_answer_length, side=\"right\")\n",
     "\n",
     "    def __getitem__(self, index):\n",
     "        x = self.questions[index]\n",
@@ -151,39 +156,43 @@
     "        x = self.text2int(x)\n",
     "        # One-hot encode x.\n",
     "        x = self.one_hot_encode(x)\n",
-    "        x = torch.tensor(x)\n",
-    "        \n",
+    "        x = torch.tensor(x).float().cuda()\n",
+    "\n",
     "        y = self.answers[index]\n",
     "        # Map text to int.\n",
     "        y = self.text2int(y)\n",
-    "        # One-hot encode y.\n",
-    "        y = self.one_hot_encode(y)\n",
-    "        y = torch.tensor(y)\n",
+    "        y = torch.tensor(y).float().cuda()\n",
     "        return x, y\n",
     "\n",
     "    def __len__(self):\n",
     "        return self.data_len\n",
-    "    \n",
+    "\n",
     "    def text2int(self, text):\n",
     "        \"\"\"\n",
     "            Convert text to an array of integers.\n",
     "        \"\"\"\n",
     "        return [self.char2int[c] for c in text]\n",
-    "    \n",
+    "\n",
+    "    def int2text(self, sequence):\n",
+    "        \"\"\"\n",
+    "            Convert an array of integers to text.\n",
+    "        \"\"\"\n",
+    "        return [self.int2char[c] for c in sequence]\n",
+    "\n",
     "    def one_hot_encode(self, sequence):\n",
     "        \"\"\"\n",
     "            Convert an array of integers to a matrix one-hot encoded.\n",
     "        \"\"\"\n",
-    "        encoded = np.zeros([self.unique_characters_length, len(sequence)], dtype=int)\n",
+    "        encoded = np.zeros([len(sequence), self.unique_characters_length], dtype=int)\n",
     "        for i, character in enumerate(sequence):\n",
-    "            encoded[character][i] = 1\n",
+    "            encoded[i][character] = 1\n",
     "        return encoded\n",
-    "    \n",
+    "\n",
     "    def one_hot_decode(self, sequence):\n",
     "        \"\"\"\n",
     "            sequence: PyTorch tensor.\n",
     "        \"\"\"\n",
-    "        return [np.argmax(x) for x in sequence.numpy().T]"
+    "        return [np.argmax(x) for x in sequence.numpy()]"
    ]
   },
   {
@@ -290,4 +299,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
\ No newline at end of file
+}
diff --git a/notebooks/network.ipynb b/notebooks/network.ipynb
index 534223c34840e6f7f12b4bf556b87f4d2f055fb8..7d055cb815c8ffd9b26de07fda2198f3eea77be6 100644
--- a/notebooks/network.ipynb
+++ b/notebooks/network.ipynb
@@ -32,41 +32,49 @@
     "\n",
     "\n",
     "class RNNModel(nn.Module):\n",
-    "    def __init__(self, input_size, output_size, hidden_dim, n_layers):\n",
+    "    def __init__(self, input_size, output_size):\n",
     "        super(RNNModel, self).__init__()\n",
     "\n",
     "        # Defining some parameters.\n",
-    "        self.hidden_dim = hidden_dim\n",
-    "        self.n_layers = n_layers\n",
+    "        self.input_size = input_size\n",
+    "        self.output_size = output_size\n",
+    "        self.hidden_dim = 32\n",
+    "        self.n_layers = 1\n",
     "\n",
     "        # region Defining the layers.\n",
     "        # RNN layer.\n",
     "        self.rnn = nn.RNN(\n",
-    "            input_size, hidden_dim, n_layers, batch_first=True, nonlinearity=\"relu\"\n",
+    "            self.input_size,\n",
+    "            self.hidden_dim,\n",
+    "            self.n_layers,\n",
+    "            batch_first=True,\n",
+    "            nonlinearity=\"relu\",\n",
     "        )\n",
     "        # Fully connected layer.\n",
-    "        self.fc = nn.Linear(hidden_dim, output_size)\n",
+    "        self.fc = nn.Linear(self.hidden_dim, self.output_size)\n",
     "        # endregion\n",
     "\n",
     "    def forward(self, x):\n",
     "        batch_size = x.size(axis=0)\n",
     "\n",
     "        # Initializing hidden state for first input using method defined below.\n",
-    "        hidden = self.init_hidden(batch_size)  # (1, 3, 12)\n",
+    "        hidden = self.init_hidden(batch_size)  # (1, 1, 32)\n",
     "\n",
     "        # Passing in the input and hidden state into the model and obtaining outputs.\n",
-    "        out, hidden = self.rnn(x, hidden)  # (3, 14, 12), (1, 3, 12)\n",
+    "        out, hidden = self.rnn(\n",
+    "            x, hidden\n",
+    "        )  # input => (3, 24, 24), (1, 1, 32) | output => (3, 24, 32), (1, 1, 32)\n",
     "\n",
     "        # Reshaping the outputs such that it can be fit into the fully connected layer.\n",
-    "        out = out.contiguous().view(-1, self.hidden_dim)  # (42, 12)\n",
-    "        out = self.fc(out)  # (42, 17)\n",
+    "        out = out.contiguous().view(-1, self.hidden_dim)  # (72, 32)\n",
+    "        out = self.fc(out)  # (72, 24)\n",
     "\n",
     "        return out, hidden\n",
     "\n",
     "    def init_hidden(self, batch_size):\n",
     "        # This method generates the first hidden state of zeros which we will use in the forward pass.\n",
     "        hidden = torch.zeros(self.n_layers, batch_size, self.hidden_dim).cuda()\n",
-    "        return hidden\n"
+    "        return hidden"
    ]
   }
  ],
diff --git a/src/main.py b/src/main.py
index d58f5834df64aedfefdc6ac5ac012f09e791a39b..e046c75e32867d417c920b643931f485e6a171b7 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,6 +1,5 @@
 import torch
 import torch.nn as nn
-from torch.utils.data.dataset import Dataset
 from torch.utils.data.sampler import SubsetRandomSampler
 import numpy as np
 from rnn import RNNModel