embed fix

7f0c12a98d1dcdc7a877f7d52df21b789a92cc6a · 2024-12-11 02:32:51 -0800 · Nel Nibcord

Files touched

M exo/inference/tinygrad/inference.py

Diff

commit 7f0c12a98d1dcdc7a877f7d52df21b789a92cc6a
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Wed Dec 11 02:32:51 2024 -0800

    embed fix
---
 exo/inference/tinygrad/inference.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index e4b76a13..e5bd3d7b 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -102,8 +102,9 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
     await self.ensure_shard(shard)
     def wrap_infer():
       x = Tensor(input_data)
-      state = self.poll_state(x, request_id)
-      out = self.model(x, **state)
+      h = self.model.embed(x)
+      state = self.poll_state(h, request_id)
+      out = self.model.forward(h, **state)
       self.states[request_id].start += x.shape[1]
       return out.realize()
     output_data = await asyncio.get_running_loop().run_in_executor(self.executor, wrap_infer)

← bd311445 Dummied up an abstact save_checkpoint · back to Exo · Proper sharding in tinygrad b1397b49 →