NUM=-2 for ants

models/efficientnet.py

@@ -100,6 +100,10 @@ class EfficientNet:
         [1, 3, (2,2), 1, 80, 192, 0.25],
         [1, 3, (2,2), 1, 192, 320, 0.25],
       ]
+    elif self.number == -2:
+      blocks_args = [
+        [1, 9, (8,8), 1, 32, 320, 0.25],
+      ]
 
     self._blocks = []
     for num_repeats, kernel_size, strides, expand_ratio, input_filters, output_filters, se_ratio in blocks_args:

test/test_mnist.py

@@ -50,8 +50,7 @@ class TestMNIST(unittest.TestCase):
     model = TinyBobNet()
     optimizer = optim.SGD(model.parameters(), lr=0.001)
     train(model, X_train, Y_train, optimizer, BS=69, steps=1)
-    for p in model.parameters():
-      p.realize()
+    for p in model.parameters(): p.realize()
 
   def test_sgd_threestep(self):
     np.random.seed(1337)
@@ -64,6 +63,7 @@ class TestMNIST(unittest.TestCase):
     model = TinyBobNet()
     optimizer = optim.Adam(model.parameters(), lr=0.001)
     train(model, X_train, Y_train, optimizer, BS=69, steps=1)
+    for p in model.parameters(): p.realize()
 
   def test_adam_threestep(self):
     np.random.seed(1337)

tinygrad/ops.py

@@ -26,6 +26,7 @@ if GRAPH:
   def save_graph_exit():
     print("saving", G)
     nx.drawing.nx_pydot.write_dot(G, '/tmp/net.dot')
+    # -Gnslimit=100 can make it finish, but you won't like results
     os.system('dot -Tsvg /tmp/net.dot -o /tmp/net.svg')
   atexit.register(save_graph_exit)