Update README.md

README.md

@@ -30,12 +30,14 @@ from PIL import Image
 
 import torch
 from transformers import AutoProcessor
-from heron.models.git_llm.git_llama import GitLlamaForCausalLM
+from heron.models.git_llm.git_gpt_neox import GitGPTNeoXForCausalLM
 
 device_id = 0
 
 # prepare a pretrained model
-model = GitLlamaForCausalLM.from_pretrained('turing-motors/heron-chat-git-ELYZA-fast-7b-v0')
+model = GitGPTNeoXForCausalLM.from_pretrained(
+    'turing-motors/heron-chat-git-ELYZA-fast-7b-v0', torch_dtype=torch.float16
+)
 model.eval()
 model.to(f"cuda:{device_id}")
 
@@ -68,7 +70,7 @@ with torch.no_grad():
     out = model.generate(**inputs, max_length=256, do_sample=False, temperature=0., eos_token_id=eos_token_id_list)
 
 # print result
-print(processor.tokenizer.batch_decode(out))
+print(processor.tokenizer.batch_decode(out)[0])
 ```