diff --git a/examples/transformer.py b/examples/transformer.py
index 9d289c8eed..5904cc9334 100755
--- a/examples/transformer.py
+++ b/examples/transformer.py
@@ -14,7 +14,7 @@ def make_dataset():
       s = i+j
       ds.append([i//10, i%10, j//10, j%10, s//100, (s//10)%10, s%10])
   random.shuffle(ds)
-  ds = np.array(ds)
+  ds = np.array(ds).astype(np.float32)
   ds_X = ds[:, 0:6]
   ds_Y = np.copy(ds[:, 1:])
   ds_X_train, ds_X_test = ds_X[0:8000], ds_X[8000:]
diff --git a/extra/training.py b/extra/training.py
index ce43fd1981..b3b3210e62 100644
--- a/extra/training.py
+++ b/extra/training.py
@@ -5,7 +5,7 @@ from tinygrad.helpers import getenv
 
 def sparse_categorical_crossentropy(out, Y):
   num_classes = out.shape[-1]
-  YY = Y.flatten()
+  YY = Y.flatten().astype(np.int32)
   y = np.zeros((YY.shape[0], num_classes), np.float32)
   # correct loss for NLL, torch NLL loss returns one per row
   y[range(y.shape[0]),YY] = -1.0*num_classes