fix test_inference_engine unittest reshape token output tensor

8a741485dfb297026e03a732be3468c5929c3c4b · 2024-11-22 14:36:20 +0400 · Alex Cheema

Files touched

M exo/inference/test_inference_engine.py

Diff

commit 8a741485dfb297026e03a732be3468c5929c3c4b
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Fri Nov 22 14:36:20 2024 +0400

    fix test_inference_engine unittest reshape token output tensor
---
 exo/inference/test_inference_engine.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/exo/inference/test_inference_engine.py b/exo/inference/test_inference_engine.py
index 690ed610..db69aebf 100644
--- a/exo/inference/test_inference_engine.py
+++ b/exo/inference/test_inference_engine.py
@@ -13,6 +13,7 @@ async def test_inference_engine(inference_engine_1: InferenceEngine, inference_e
   prompt = "In a single word only, what is the last name of the current president of the USA?"
   resp_full = await inference_engine_1.infer_prompt("A", shard=Shard(model_id=model_id, start_layer=0, end_layer=n_layers - 1, n_layers=n_layers), prompt=prompt)
   token_full = await inference_engine_1.sample(resp_full)
+  token_full = token_full.reshape(1, -1)
   next_resp_full = await inference_engine_1.infer_tensor(
     "A",
     shard=Shard(model_id=model_id, start_layer=0, end_layer=n_layers - 1, n_layers=n_layers),
@@ -27,6 +28,7 @@ async def test_inference_engine(inference_engine_1: InferenceEngine, inference_e
     input_data=resp1,
   )
   tokens2 = await inference_engine_1.sample(resp2)
+  tokens2 = tokens2.reshape(1, -1)
   resp3 = await inference_engine_1.infer_tensor(
     "B",
     shard=Shard(model_id=model_id, start_layer=0, end_layer=pp, n_layers=n_layers),

← e3ec9eaa Fixed GRPC issues · back to Exo · run tests with --disable-tui for faster tests d3505e03 →