Spaces:

poltextlab
/

babelmachine

Running

poltextlab commited on May 23

Commit

ac10520

•

1 Parent(s): 7658607

offload

Files changed (1) hide show

utils.py CHANGED Viewed

@@ -23,7 +23,7 @@ tokenizers = ["xlm-roberta-large"]
 def download_hf_models():
     for model_id in models:
-        AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto",
                                                                    token=HF_TOKEN)
     for tokenizer_id in tokenizers:
         AutoTokenizer.from_pretrained(tokenizer_id)

 def download_hf_models():
     for model_id in models:
+        AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", offload_folder="offload",
                                                                    token=HF_TOKEN)
     for tokenizer_id in tokenizers:
         AutoTokenizer.from_pretrained(tokenizer_id)