Spaces:

codeparrot
/

codeparrot-generation

Build error

lvwerra HF Staff commited on Apr 21, 2022

Commit

d5f0a3f

1 Parent(s): 7dea8d5

low mem loading of model

Files changed (1) hide show

app.py CHANGED Viewed

@@ -11,7 +11,7 @@ def load_tokenizer(model_ckpt):
 @st.cache(allow_output_mutation=True)
 def load_model(model_ckpt):
-    model = AutoModelForCausalLM.from_pretrained(model_ckpt)
     return model
 @st.cache()

 @st.cache(allow_output_mutation=True)
 def load_model(model_ckpt):
+    model = AutoModelForCausalLM.from_pretrained(model_ckpt, low_cpu_mem_usage=True)
     return model
 @st.cache()