From 9594b2db1c9ca1cbf4350805c1860e054eb1ea4a Mon Sep 17 00:00:00 2001
From: Gnome Ann <>
Date: Mon, 17 Jan 2022 13:17:20 -0500
Subject: [PATCH] Fix soft prompt length calculation in `calcsubmitbudget()`

In TPU instances, `vars.sp.shape[0]` is not always the actual number of
tokens in the soft prompt. We have to use `vars.sp_length` to get an
accurate token count.
---
 aiserver.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/aiserver.py b/aiserver.py
index d1a7efc8..00e96e32 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -2494,7 +2494,7 @@ def calcsubmitbudget(actionlen, winfo, mem, anotetxt, actions, submission=None,
     anotetkns    = []  # Placeholder for Author's Note tokens
     lnanote      = 0   # Placeholder for Author's Note length
 
-    lnsp = vars.sp.shape[0] if vars.sp is not None else 0
+    lnsp = vars.sp_length
 
     if("tokenizer" not in globals()):
         from transformers import GPT2TokenizerFast