diff --git a/test/mnist.py b/test/mnist.py
index 686b6fcfc9..e950368b8f 100644
--- a/test/mnist.py
+++ b/test/mnist.py
@@ -1,8 +1,8 @@
 #!/usr/bin/env python
 import numpy as np
 from tinygrad.tensor import Tensor
-from tinygrad.nn import layer_init, SGD
 from tinygrad.utils import fetch_mnist
+import tinygrad.optim as optim
 
 from tqdm import trange
 
@@ -12,6 +12,10 @@ X_train, Y_train, X_test, Y_test = fetch_mnist()
 
 # train a model
 
+def layer_init(m, h):
+  ret = np.random.uniform(-1., 1., size=(m,h))/np.sqrt(m*h)
+  return ret.astype(np.float32)
+
 class TinyBobNet:
   def __init__(self):
     self.l1 = Tensor(layer_init(784, 128))
@@ -22,9 +26,8 @@ class TinyBobNet:
 
 # optimizer
 
-
 model = TinyBobNet()
-optim = SGD([model.l1, model.l2], lr=0.01)
+optim = optim.SGD([model.l1, model.l2], lr=0.01)
 
 BS = 128
 losses, accuracies = [], []
diff --git a/tinygrad/__init__.py b/tinygrad/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tinygrad/nn.py b/tinygrad/optim.py
similarity index 55%
rename from tinygrad/nn.py
rename to tinygrad/optim.py
index e3f6e95f62..e8adbb5776 100644
--- a/tinygrad/nn.py
+++ b/tinygrad/optim.py
@@ -1,9 +1,3 @@
-import numpy as np
-
-def layer_init(m, h):
-  ret = np.random.uniform(-1., 1., size=(m,h))/np.sqrt(m*h)
-  return ret.astype(np.float32)
-
 class SGD:
   def __init__(self, tensors, lr):
     self.tensors = tensors