diff --git a/README.md b/README.md
index 5266f7a..fad9f4a 100644
--- a/README.md
+++ b/README.md
@@ -74,7 +74,7 @@ import random
 random.seed(1)
 
 BATCH_SIZE = 32
-device = "cpu"
+device = "cuda" #cpu
 epochs = 10
 
 transform = transforms.Sequential(
diff --git a/examples/train.ipynb b/examples/train.ipynb
index f1e51a5..c53ea86 100644
--- a/examples/train.ipynb
+++ b/examples/train.ipynb
@@ -99,7 +99,7 @@
    "outputs": [],
    "source": [
     "BATCH_SIZE = 32\n",
-    "device = \"cpu\"\n",
+    "device = \"cuda\" #cpu\n",
     "epochs = 10\n",
     "\n",
     "transform = transforms.Sequential(\n",
diff --git a/norch/nn/functional.py b/norch/nn/functional.py
index d63e3d2..29fb68d 100644
--- a/norch/nn/functional.py
+++ b/norch/nn/functional.py
@@ -29,6 +29,4 @@ def one_hot_encode(x, num_classes):
         target_idx = int(x[i])
         one_hot[i][target_idx] = 1
 
-    
-
     return norch.Tensor(one_hot)
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 06c8016..36bff91 100644
--- a/setup.py
+++ b/setup.py
@@ -20,7 +20,7 @@ class CustomInstall(install):
 
 setuptools.setup(
     name = "norch",
-    version = "0.0.3",
+    version = "0.0.4",
     author = "Lucas de Lima",
     author_email = "nogueiralucasdelima@gmail.com",
     description = "A deep learning framework",
diff --git a/test.py b/test.py
deleted file mode 100644
index 9297870..0000000
--- a/test.py
+++ /dev/null
@@ -1,27 +0,0 @@
-import norch
-import norch.nn as nn
-import norch.optim as optim
-from norch.utils.data.dataloader import Dataloader
-from norch.norchvision import transforms
-import numpy as np
-import matplotlib.pyplot as plt
-import random
-random.seed(1)
-
-class MyModel(nn.Module):
-    def __init__(self):
-        super(MyModel, self).__init__()
-        self.fc1 = nn.Linear(784, 30)
-        self.sigmoid1 = nn.Sigmoid()
-        self.fc2 = nn.Linear(30, 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
-
-model = MyModel().to("cuda")
\ No newline at end of file