Fix TPU

2022-06-06 21:39:17 -04:00 · 2022-06-06 21:39:17 -04:00 · 1b35b55d86
parent ae1aed0916
commit 1b35b55d86
1 changed files with 1 additions and 0 deletions
--- 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(