diff --git a/norch/nn/__pycache__/loss.cpython-38.pyc b/norch/nn/__pycache__/loss.cpython-38.pyc
index b8133b9..9d254eb 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/functional.py b/norch/nn/functional.py
index 9493591..78b1f53 100644
--- a/norch/nn/functional.py
+++ b/norch/nn/functional.py
@@ -27,7 +27,4 @@ def one_hot_encode(x, num_classes):
         target_idx = int(x.tensor.contents.data[i])
         one_hot[i][target_idx] = 1
 
-    if x.numel < 2:
-        one_hot = one_hot[0]
-
     return norch.Tensor(one_hot)
\ No newline at end of file
diff --git a/norch/nn/loss.py b/norch/nn/loss.py
index 8d7b0f9..b0e7898 100644
--- a/norch/nn/loss.py
+++ b/norch/nn/loss.py
@@ -45,6 +45,7 @@ class CrossEntropyLoss(Loss):
                 target = norch.one_hot_encode(target, num_classes)
                 
                 logits = norch.softmax(input, dim=0)
+                target = target.reshape(logits.shape)
                 cost = -(logits.log() * target).sum()
                 
             else:
@@ -52,6 +53,7 @@ class CrossEntropyLoss(Loss):
                 assert target.shape == input.shape, \
                     "Input and target shape does not match: {} and {}".format(input.shape, target.shape)
                 logits = norch.softmax(input, dim=0)
+                target = target.reshape(logits.shape)
                 cost = -(logits.log() * target).sum()
 
 
@@ -65,6 +67,7 @@ class CrossEntropyLoss(Loss):
                 
                 batch_size = input.shape[0]
                 logits = norch.softmax(input, dim=1)
+                target = target.reshape(logits.shape)
                 cost = -(logits.log() * target).sum() / batch_size
 
             else:
@@ -74,6 +77,7 @@ class CrossEntropyLoss(Loss):
                 
                 batch_size = input.shape[0]
                 logits = norch.softmax(input, dim=1)
+                target = target.reshape(logits.shape)
                 cost = -(logits.log() * target).sum() / batch_size
 
         if input.requires_grad:
diff --git a/test.py b/test.py
index 7ad7e2d..23c7928 100644
--- a/test.py
+++ b/test.py
@@ -23,13 +23,15 @@ class MyModel(nn.Module):
     def __init__(self):
         super(MyModel, self).__init__()
         self.fc1 = nn.Linear(784, 5)
-        self.sigmoid = nn.Sigmoid()
+        self.sigmoid1 = nn.Sigmoid()
         self.fc2 = nn.Linear(5, 10)
+        self.sigmoid2 = nn.Sigmoid()
 
     def forward(self, x):
         out = self.fc1(x)
-        out = self.sigmoid(out)
+        out = self.sigmoid1(out)
         out = self.fc2(out)
+        out = self.sigmoid2(out)
         
         return out
 
@@ -49,10 +51,11 @@ for epoch in range(epochs):
         target = target
 
         x = x.to(device)
-        target = target.to(device).unsqueeze(-1)
+        target = target.to(device)
 
-        outputs = model(x)
-        print(outputs.shape, target.shape, x.shape)
+        outputs = model(x).squeeze(-1)
+        print(x.shape, outputs.shape, target.shape)
+        
         loss = criterion(outputs, target)
         
         optimizer.zero_grad()
@@ -70,6 +73,7 @@ for epoch in range(epochs):
         print('f1 depois', model.fc1.bias)
         print('f2 depois', model.fc2.bias)
         print('\n\n')
+        exit()
 
     print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}')
     loss_list.append(loss[0])
\ No newline at end of file