Update README.md

README.md

@@ -21,6 +21,7 @@ use oig-sql dataset and support more complex sql parse
 # How to Use
 
 ```python
+import torch
 from transformers import AutoTokenizer, BartForConditionalGeneration
 
 device = torch.device('cuda:0')
@@ -30,7 +31,7 @@ model = BartForConditionalGeneration.from_pretrained("LarkAI/bart_large_nl2sql")
 
 text = "question: get people name with age less 25 table: id, name, age"
 inputs = tokenizer([text], max_length=1024, return_tensors="pt")
-output_ids = model.generate(inputs["input_ids"].to(self.device), num_beams=self.beams, max_length=128, min_length=8)
+output_ids = model.generate(inputs["input_ids"].to(device), num_beams=self.beams, max_length=128, min_length=8)
 response_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
 # SELECT name FROM table WHERE age < 25
 ```