Add a --torch-compile flag for transformers

This commit is contained in:
oobabooga 2025-01-05 05:45:12 -08:00
parent 11af199aff
commit c0f600c887
5 changed files with 9 additions and 2 deletions

View file

@ -254,6 +254,9 @@ def huggingface_loader(model_name):
print()
model = LoaderClass.from_pretrained(path_to_model, **params)
if shared.args.torch_compile:
model = torch.compile(model)
return model