Do we need casting here?

59af2dd592963e9ba464eb856867c08065ba475f · 2024-12-08 02:50:12 -0800 · Nel Nibcord

Files touched

M exo/inference/mlx/sharded_inference_engine.py

Diff

commit 59af2dd592963e9ba464eb856867c08065ba475f
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Sun Dec 8 02:50:12 2024 -0800

    Do we need casting here?
---
 exo/inference/mlx/sharded_inference_engine.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 3ba8370e..dc1e17e8 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -81,7 +81,7 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
     await self.ensure_shard(shard)
     loop = asyncio.get_running_loop()
     state = await self.poll_state(request_id)
-    x = mx.array(input_data).astype(mx.int64) if self.shard.is_first_layer() else mx.array(input_data)
+    x = mx.array(input_data)
     output_data: np.ndarray = np.array(await loop.run_in_executor(self.executor, lambda: self.model(x, **state)))
     return output_data
 
@@ -90,7 +90,7 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
     await self.save_session('loss', loss_fns[loss])
     loop = asyncio.get_running_loop()
     #print(f"evaluate in <- {inputs}")
-    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)
     score = await loop.run_in_executor(self.executor, self.session['loss'], self.model, x, y, l)

← b22c21ac Some session method cleanup · back to Exo · Updated node refs 763fbf84 →