Fix for loading models on CPU only that don't support breakmodel

This commit is contained in:
ebolam 2022-09-27 18:29:32 -04:00
parent edd50fc809
commit 0311cc215e
1 changed files with 2 additions and 0 deletions

View File

@ -1514,8 +1514,10 @@ def get_model_info(model, directory=""):
pass
else:
layer_count = get_layer_count(model, directory=directory)
print("Layer Count: {}".format(layer_count))
if layer_count is None:
breakmodel = False
gpu = True
else:
breakmodel = True
if model in ["NeoCustom", "GPT2Custom"]: