From 1b35b55d864646006e262fbc34900f91391bf969 Mon Sep 17 00:00:00 2001
From: ebolam <ebolam@gmail.com>
Date: Mon, 6 Jun 2022 21:39:17 -0400
Subject: [PATCH] Fix TPU

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

diff --git a/aiserver.py b/aiserver.py
index f6fcab7b..5bcde1a8 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -5611,6 +5611,7 @@ def final_startup():
 
     # Precompile TPU backend if required
     if(vars.use_colab_tpu or vars.model in ("TPUMeshTransformerGPTJ", "TPUMeshTransformerGPTNeoX")):
+        import tpu_mtj_backend
         soft_tokens = tpumtjgetsofttokens()
         if(vars.dynamicscan or (not vars.nogenmod and vars.has_genmod)):
             threading.Thread(