diff --git a/examples/train.ipynb b/examples/train.ipynb
index b29a29d..79e90fb 100644
--- a/examples/train.ipynb
+++ b/examples/train.ipynb
@@ -20005,17 +20005,20 @@
     "class MyModel(nn.Module):\n",
     "    def __init__(self):\n",
     "        super(MyModel, self).__init__()\n",
-    "        self.fc1 = nn.Linear(1, 10)\n",
-    "        self.sigmoid = nn.Sigmoid()\n",
-    "        self.fc2 = nn.Linear(10, 1)\n",
+    "        self.fc1 = nn.Linear(784, 5)\n",
+    "        self.sigmoid1 = nn.Sigmoid()\n",
+    "        self.fc2 = nn.Linear(5, 10)\n",
+    "        self.sigmoid2 = nn.Sigmoid()\n",
     "\n",
     "    def forward(self, x):\n",
     "        out = self.fc1(x)\n",
-    "        out = self.sigmoid(out)\n",
+    "        out = self.sigmoid1(out)\n",
     "        out = self.fc2(out)\n",
+    "        out = self.sigmoid2(out)\n",
     "        \n",
     "        return out\n",
     "\n",
+    "\n",
     "device = \"cpu\"\n",
     "epochs = 10\n",
     "\n",
@@ -20024,15 +20027,9 @@
     "optimizer = optim.SGD(model.parameters(), lr=0.001)\n",
     "loss_list = []\n",
     "\n",
-    "x_values = [0. ,  0.4,  0.8,  1.2,  1.6,  2. ,  2.4,  2.8,  3.2,  3.6,  4. ,\n",
-    "        4.4,  4.8,  5.2,  5.6,  6. ,  6.4,  6.8,  7.2,  7.6,  8. ,  8.4,\n",
-    "        8.8,  9.2,  9.6, 10. , 10.4, 10.8, 11.2, 11.6, 12. , 12.4, 12.8,\n",
-    "       13.2, 13.6, 14. , 14.4, 14.8, 15.2, 15.6, 16. , 16.4, 16.8, 17.2,\n",
-    "       17.6, 18. , 18.4, 18.8, 19.2, 19.6, 20.]\n",
+    "x_values = [1 for i in range(784)]\n",
     "\n",
-    "y_true = []\n",
-    "for x in x_values:\n",
-    "    y_true.append(math.pow(math.sin(x), 2))\n",
+    "y_true = [1 for i in range(784)]\n",
     "\n",
     "batch_size = 100\n",
     "\n",
@@ -20041,7 +20038,8 @@
     "    for i, (x, target) in enumerate(zip(x_values, y_true)):\n",
     "        x = norch.Tensor([[x] for _ in range(batch_size)]).T\n",
     "        target = norch.Tensor([[target] for _ in range(batch_size)]).T\n",
-    "\n",
+    "        print(x.shape)\n",
+    "        print(target.shape)\n",
     "        x = x.to(device)\n",
     "        target = target.to(device)\n",
     "\n",
diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc
index bcc1fe0..c31fb5d 100644
Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ
diff --git a/norch/nn/__pycache__/loss.cpython-38.pyc b/norch/nn/__pycache__/loss.cpython-38.pyc
index 9d254eb..40e2a17 100644
Binary files a/norch/nn/__pycache__/loss.cpython-38.pyc and b/norch/nn/__pycache__/loss.cpython-38.pyc differ
diff --git a/norch/nn/loss.py b/norch/nn/loss.py
index b0e7898..bb5ddd8 100644
--- a/norch/nn/loss.py
+++ b/norch/nn/loss.py
@@ -58,6 +58,8 @@ class CrossEntropyLoss(Loss):
 
 
         elif input.ndim == 2:
+            if target.ndim > 1:
+                target = target.squeeze(-1)
             # batched 
             if target.ndim == 1:
                 # target -> Ground truth class indices:
diff --git a/norch/tensor.py b/norch/tensor.py
index 1121066..a2f0cb5 100644
--- a/norch/tensor.py
+++ b/norch/tensor.py
@@ -178,7 +178,8 @@ class Tensor:
             
             # Only squeeze the specified dimension if its size is 1
             if self.shape[dim] != 1:
-                raise ValueError("Dimension {0} does not have size 1 and cannot be squeezed".format(dim))
+                return self
+                #raise ValueError("Dimension {0} does not have size 1 and cannot be squeezed".format(dim))
             
             # Create the new shape without the specified dimension
             new_shape = self.shape[:dim] + self.shape[dim+1:]
diff --git a/test.py b/test.py
index 23c7928..9468eb0 100644
--- a/test.py
+++ b/test.py
@@ -1,4 +1,63 @@
 import norch
+import norch.nn as nn
+import norch.optim as optim
+import random
+import math
+
+"""random.seed(1)
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super(MyModel, self).__init__()
+        self.fc1 = nn.Linear(784, 5)
+        self.sigmoid1 = nn.Sigmoid()
+        self.fc2 = nn.Linear(5, 10)
+        self.sigmoid2 = nn.Sigmoid()
+
+    def forward(self, x):
+        out = self.fc1(x)
+        out = self.sigmoid1(out)
+        out = self.fc2(out)
+        out = self.sigmoid2(out)
+        
+        return out
+
+
+device = "cpu"
+epochs = 1000
+
+model = MyModel().to(device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.SGD(model.parameters(), lr=0.1)
+loss_list = []
+
+x_values = [1 for i in range(784)]
+
+y_true = [1]
+
+batch_size = 10
+
+
+for epoch in range(epochs):
+
+    x = norch.Tensor([x_values for _ in range(batch_size)])
+    target = norch.Tensor([y_true for _ in range(batch_size)])
+    x = x.to(device).unsqueeze(-1)
+    target = target.to(device)
+
+    outputs = model(x).squeeze(-1)
+
+    loss = criterion(outputs, target)
+    optimizer.zero_grad()
+    loss.backward()
+        
+
+    optimizer.step()
+    print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}')
+    loss_list.append(loss[0])
+
+"""
+import norch
 from norch.utils.data.dataloader import Dataloader
 import norch
 import norch.nn as nn
@@ -15,7 +74,7 @@ target_transform = lambda x: to_tensor(x)
 train_data, test_data = norch.datasets.MNIST.splits(transform=transform, target_transform=target_transform)
 sample, _ = train_data[0]
 
-BATCH_SIZE = 100
+BATCH_SIZE = 1
 
 train_loader = Dataloader(train_data, batch_size = BATCH_SIZE)
 
@@ -40,7 +99,7 @@ epochs = 10
 
 model = MyModel().to(device)
 criterion = nn.CrossEntropyLoss()
-optimizer = optim.SGD(model.parameters(), lr=0.001)
+optimizer = optim.SGD(model.parameters(), lr=0.01)
 loss_list = []
 
 for epoch in range(epochs):    
@@ -53,7 +112,9 @@ for epoch in range(epochs):
         x = x.to(device)
         target = target.to(device)
 
+        
         outputs = model(x).squeeze(-1)
+        print(outputs.shape)
         print(x.shape, outputs.shape, target.shape)
         
         loss = criterion(outputs, target)