Spaces:

numind
/

NuExtract

Running on L4

Alexandre-Numind commited on Jun 20

Commit

4eb458a

•

1 Parent(s): bb9a67b

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -306,9 +306,9 @@ def highlight_words(input_text, json_output):
 model = AutoModelForCausalLM.from_pretrained(
         "numind/NuExtract",
         trust_remote_code=True,
-        device_map="cuda",
         )
 tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract")
 tokenizer.eos = tokenizer("<|end-output|>")

 model = AutoModelForCausalLM.from_pretrained(
         "numind/NuExtract",
         trust_remote_code=True,
         )
+model.to("cuda")
 tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract")
 tokenizer.eos = tokenizer("<|end-output|>")