From 32ed128598b8d3fea08cf2402242e4d048e2dd91 Mon Sep 17 00:00:00 2001
From: Nishant Rajadhyaksha <79035403+nishant42491@users.noreply.github.com>
Date: Sun, 13 Apr 2025 16:34:20 -0700
Subject: [PATCH] fixing transformer training bug (#9877)

---
 examples/transformer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/transformer.py b/examples/transformer.py
index 375f4ca907..40082db71c 100755
--- a/examples/transformer.py
+++ b/examples/transformer.py
@@ -17,7 +17,7 @@ def make_dataset():
   random.shuffle(ds)
   ds = np.array(ds).astype(np.float32)
   ds_X = ds[:, 0:6]
-  ds_Y = np.copy(ds[:, 1:])
+  ds_Y = np.copy(ds[:, 1:]).astype(np.int32)
   ds_X_train, ds_X_test = ds_X[0:8000], ds_X[8000:]
   ds_Y_train, ds_Y_test = ds_Y[0:8000], ds_Y[8000:]
   return ds_X_train, ds_Y_train, ds_X_test, ds_Y_test