From 1d32b090a9388feedc357d702d19b93deea54e01 Mon Sep 17 00:00:00 2001
From: somebody <onesome01@protonmail.com>
Date: Thu, 22 Sep 2022 21:08:36 -0500
Subject: [PATCH] MTJ Gen Work

---
 aiserver.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/aiserver.py b/aiserver.py
index 7d9d8a17..e79c0272 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -5053,6 +5053,8 @@ def tpu_raw_generate(
     max_length: int,
     batch_count: int,
 ):
+
+    prompt_tokens = prompt_tokens[0]
     # Mostly lifted from apiactionsubmit_tpumtjgenerate
     soft_tokens = tpumtjgetsofttokens()
     __debug("we are generating with", prompt_tokens, "batch", batch_count, "soft tokens", soft_tokens)