Update README.md

README.md

@@ -42,6 +42,7 @@ To use the model:
 ```python
 import json
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 
 def predict_NuExtract(model,tokenizer,text, schema,example = ["","",""]):
@@ -61,7 +62,7 @@ def predict_NuExtract(model,tokenizer,text, schema,example = ["","",""]):
 model = AutoModelForCausalLM.from_pretrained("numind/NuExtract", trust_remote_code=True,torch_dtype=torch.bfloat16)
 tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract", trust_remote_code=True)
 
-#model.to("cuda")
+model.to("cuda")
 
 model.eval()
 
@@ -91,6 +92,6 @@ schema = """{
 }"""
 
 prediction = predict_NuExtract(model,tokenizer,text, schema,example = ["","",""])
-
+print(prediction)
 
 ```