Change default `model.save_pretrained` shard size to 500 MiB

This commit is contained in:
Gnome Ann 2022-05-10 22:04:25 -04:00
parent a388c63023
commit 937d9ee06a
1 changed files with 1 additions and 1 deletions

View File

@ -1506,7 +1506,7 @@ if(not vars.use_colab_tpu and vars.model not in ["InferKit", "Colab", "OAI", "Go
if not args.colab or args.savemodel:
import shutil
model = model.half()
model.save_pretrained("models/{}".format(vars.model.replace('/', '_')))
model.save_pretrained("models/{}".format(vars.model.replace('/', '_')), max_shard_size="500MiB")
tokenizer.save_pretrained("models/{}".format(vars.model.replace('/', '_')))
shutil.rmtree("cache/")