Update README.md

README.md

@@ -106,7 +106,7 @@ from transformers import (
     GenerationConfig
 )
 from peft import PeftModel
-
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_path = 'baichuan-inc/Baichuan2-13B-Chat'
 lora_path = 'zjunlp/baichuan2-13b-iepile-lora'
 config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
@@ -130,7 +130,7 @@ model.eval()
 sintruct = "{\"instruction\": \"You are an expert in named entity recognition. Please extract entities that match the schema definition from the input. Return an empty list if the entity type does not exist. Please respond in the format of a JSON string.\", \"schema\": [\"person\", \"organization\", \"else\", \"location\"], \"input\": \"284 Robert Allenby ( Australia ) 69 71 71 73 , Miguel Angel Martin ( Spain ) 75 70 71 68 ( Allenby won at first play-off hole )\"}"
 sintruct = '<reserved_106>' + sintruct + '<reserved_107>'
 
-input_ids = tokenizer.encode(sintruct, return_tensors="pt")
+input_ids = tokenizer.encode(sintruct, return_tensors="pt").to(device)
 input_length = input_ids.size(1)
 generation_output = model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_length=512, max_new_tokens=256, return_dict_in_generate=True))
 generation_output = generation_output.sequences[0]