From ae837334316f49d2c85c7c60ccfd2c8dc2a703fe Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Tue, 9 Jan 2024 19:25:37 -0800
Subject: [PATCH] hotfix: examples/transformer.py

---
 examples/transformer.py | 5 +++--
 extra/training.py       | 5 +++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/examples/transformer.py b/examples/transformer.py
index af20a44904..56c573ba0c 100755
--- a/examples/transformer.py
+++ b/examples/transformer.py
@@ -28,7 +28,8 @@ if __name__ == "__main__":
   lr = 0.003
   for i in range(10):
     optim = Adam(get_parameters(model), lr=lr)
-    train(model, X_train, Y_train, optim, 50, BS=64)
+    # TODO: BUG! why doesn't the JIT work here?
+    train(model, X_train, Y_train, optim, 50, BS=64, allow_jit=False)
     acc, Y_test_preds = evaluate(model, X_test, Y_test, num_classes=10, return_predict=True)
     lr /= 1.2
     print(f'reducing lr to {lr:.4f}')
@@ -37,6 +38,6 @@ if __name__ == "__main__":
       for k in range(len(Y_test_preds)):
         if (Y_test_preds[k] != Y_test[k]).any():
           wrong+=1
-          a,b,c,x = X_test[k,:2], X_test[k,2:4], Y_test[k,-3:], Y_test_preds[k,-3:]
+          a,b,c,x = X_test[k,:2].astype(np.int32), X_test[k,2:4].astype(np.int32), Y_test[k,-3:].astype(np.int32), Y_test_preds[k,-3:].astype(np.int32)
           print(f'{a[0]}{a[1]} + {b[0]}{b[1]} = {x[0]}{x[1]}{x[2]} (correct: {c[0]}{c[1]}{c[2]})')
       print(f'Wrong predictions: {wrong}, acc = {acc:.4f}')
diff --git a/extra/training.py b/extra/training.py
index c127ae6f3a..8862fec963 100644
--- a/extra/training.py
+++ b/extra/training.py
@@ -6,9 +6,8 @@ from tinygrad.jit import TinyJit
 
 
 def train(model, X_train, Y_train, optim, steps, BS=128, lossfn=lambda out,y: out.sparse_categorical_crossentropy(y),
-        transform=lambda x: x, target_transform=lambda x: x, noloss=False):
+        transform=lambda x: x, target_transform=lambda x: x, noloss=False, allow_jit=True):
 
-  @TinyJit
   def train_step(x, y):
     # network
     out = model.forward(x) if hasattr(model, 'forward') else model(x)
@@ -22,6 +21,8 @@ def train(model, X_train, Y_train, optim, steps, BS=128, lossfn=lambda out,y: ou
     accuracy = (cat == y).mean()
     return loss.realize(), accuracy.realize()
 
+  if allow_jit: train_step = TinyJit(train_step)
+
   with Tensor.train():
     losses, accuracies = [], []
     for i in (t := trange(steps, disable=CI)):