From 216eb235e5dad207aaaeecb709210abdd825d96f Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Tue, 9 Apr 2024 19:30:03 -0700
Subject: [PATCH] hotfix: cast mnist to float

---
 examples/beautiful_mnist.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/examples/beautiful_mnist.py b/examples/beautiful_mnist.py
index 4c56b32ddf..7c55d204da 100644
--- a/examples/beautiful_mnist.py
+++ b/examples/beautiful_mnist.py
@@ -21,6 +21,9 @@ class Model:
 if __name__ == "__main__":
   X_train, Y_train, X_test, Y_test = mnist()
 
+  # TODO: remove this when HIP is fixed
+  X_train, X_test = X_train.float(), X_test.float()
+
   model = Model()
   opt = nn.optim.Adam(nn.state.get_parameters(model))