From 45bfde8d5d41f1104b8783e1f21e17749e30a613 Mon Sep 17 00:00:00 2001
From: Gnome Ann <>
Date: Mon, 17 Jan 2022 13:19:57 -0500
Subject: [PATCH] `generated_cols` needs to be set properly by TPU static
 backend

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

diff --git a/aiserver.py b/aiserver.py
index 00e96e32..d25dfe11 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -3064,6 +3064,7 @@ def tpumtjgenerate(txt, minimum, maximum, found_entries=None):
             past = genout
             for i in range(vars.numseqs):
                 vars.lua_koboldbridge.generated[i+1] = vars.lua_state.table(*genout[i].tolist())
+            vars.lua_koboldbridge.generated_cols = vars.generated_tkns = genout.shape[-1]
 
     except Exception as e:
         if(issubclass(type(e), lupa.LuaError)):