Fix TPU
This commit is contained in:
parent
ae1aed0916
commit
1b35b55d86
|
@ -5611,6 +5611,7 @@ def final_startup():
|
||||||
|
|
||||||
# Precompile TPU backend if required
|
# Precompile TPU backend if required
|
||||||
if(vars.use_colab_tpu or vars.model in ("TPUMeshTransformerGPTJ", "TPUMeshTransformerGPTNeoX")):
|
if(vars.use_colab_tpu or vars.model in ("TPUMeshTransformerGPTJ", "TPUMeshTransformerGPTNeoX")):
|
||||||
|
import tpu_mtj_backend
|
||||||
soft_tokens = tpumtjgetsofttokens()
|
soft_tokens = tpumtjgetsofttokens()
|
||||||
if(vars.dynamicscan or (not vars.nogenmod and vars.has_genmod)):
|
if(vars.dynamicscan or (not vars.nogenmod and vars.has_genmod)):
|
||||||
threading.Thread(
|
threading.Thread(
|
||||||
|
|
Loading…
Reference in New Issue