diff --git a/main.py b/main.py
index 43cbb97..e7cc6e1 100644
--- a/main.py
+++ b/main.py
@@ -28,8 +28,8 @@
 
 net.load("10epoch_weights.pkl")
 
-#net.train(training_data[:100], training_labels[:100], 32, 3, 'weights.pkl')
+net.train(training_data, training_labels, 32, 10, 'weights.pkl')
 
-net.test(testing_data[:500], testing_labels[:500])
+net.test(testing_data[:50], testing_labels[:50])
 
 # save_vanilla_gradient(net, training_data[:25], training_labels[:25], 5)
diff --git a/model/layers.py b/model/layers.py
index 2a69dc5..441888b 100644
--- a/model/layers.py
+++ b/model/layers.py
@@ -147,7 +147,7 @@ def backward(self, dy):
 
     def parameters(self):
         return
-    
+
     def load(self, weights, bias):
         return
 
@@ -231,7 +231,7 @@ def __init__(self):
         self.out = None
 
     def forward(self, inputs):
-        exp = np.exp(inputs, dtype=np.float)
+        exp = np.exp(inputs, dtype=np.float64)
         self.out = exp/np.sum(exp)
         return self.out
 
@@ -240,6 +240,6 @@ def backward(self, dy):
 
     def parameters(self):
         return
-    
+
     def load(self, weights, bias):
         return
diff --git a/model/loss.py b/model/loss.py
index 7248a3a..1996c99 100644
--- a/model/loss.py
+++ b/model/loss.py
@@ -3,6 +3,6 @@
 
 def cross_entropy(inputs, labels):
     out_num = labels.shape[0]
-    probability = np.sum(labels.reshape(1, out_num) * inputs)
-    loss = -np.log(probability)
+    probability = np.sum(labels.reshape(1, out_num) * inputs, dtype=np.float64)
+    loss = -np.log(probability, dtype=np.float64)
     return loss