Missed one

cc66a0b7822d08f9a8ff069fdd8422bd9501cbed · 2024-12-10 13:15:19 -0800 · Nel Nibcord

Files touched

M exo/inference/mlx/sharded_inference_engine.py

Diff

commit cc66a0b7822d08f9a8ff069fdd8422bd9501cbed
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Tue Dec 10 13:15:19 2024 -0800

    Missed one
---
 exo/inference/mlx/sharded_inference_engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index dc1e17e8..a7386fb6 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -120,7 +120,7 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
       mx.eval(self.model.parameters(), self.session['opt'].state, lval)
       return lval, gradlayers
 
-    x = mx.array(inputs).astype(mx.int64) if self.shard.is_first_layer() else mx.array(inputs)
+    x = mx.array(inputs)
     y = mx.array(targets)
     l = mx.array(lengths)

← 124a0338 Slightly simplified waiting for outstanding requests · back to Exo · Dummied up an abstact save_checkpoint bd311445 →