diff --git a/assignment 5/iml_assignmnet5_unsolved.ipynb b/assignment 5/iml_assignmnet5_unsolved.ipynb
index bd06111..b115b6d 100644
--- a/assignment 5/iml_assignmnet5_unsolved.ipynb	
+++ b/assignment 5/iml_assignmnet5_unsolved.ipynb	
@@ -75,15 +75,28 @@
    "outputs": [],
    "source": [
     "# Define the perceptron algorithm\n",
-    "class Perceptron:\n",
-    "    def __init__(self, learning_rate=0.01, n_iters=1000):\n",
+    "class MultiClassPerceptron:\n",
+    "    def __init__(self, input_dim, output_dim, lr=0.01, epochs=1000):\n",
+    "        self.W = np.random.randn(input_dim, output_dim)\n",
+    "        self.b = np.zeros((1, output_dim))\n",
+    "        self.lr = lr\n",
+    "        self.epochs = epochs\n",
+    "\n",
+    "    def forward(self, X):\n",
+    "        # ToDo: implement the forward() function\n",
     "        pass\n",
     "\n",
-    "    # define the fit function to train the model\n",
+    "    def backward(self, X, y):\n",
+    "       # ToDo: implement the backward() function\n",
+    "        pass\n",
     "\n",
-    "    # define the predict function to predict labels\n",
+    "    def fit(self, X, y):\n",
+    "        for epoch in range(self.epochs):\n",
+    "            self.forward(X)\n",
+    "            self.backward(X, y)\n",
     "\n",
-    "    def _unit_step_func(self, x):\n",
+    "    def predict(self, X):\n",
+    "        # ToDo: implement the predict() function\n",
     "        pass"
    ],
    "metadata": {
@@ -105,7 +118,7 @@
    "outputs": [],
    "source": [
     "# Train the model\n",
-    "p = Perceptron(learning_rate=0.01, n_iters=1000)\n",
+    "p = MultiClassPerceptron(input_dim=X_train.shape[1], output_dim=3, lr=0.01, epochs=1000)\n",
     "p.fit(X_train, y_train)\n",
     "predictions_train = p.predict(X_train)\n",
     "predictions = p.predict(X_test)"