Update README.md

README.md

@@ -3,6 +3,7 @@ kakao brain에서 공개한 kogpt 6b model('kakaobrain/kogpt')을 fp16으로 저
 ### 카카오브레인 모델을 fp16으로 로드하는 방법
 
 ```
+import torch
 from transformers import GPTJForCausalLM
 
 model = GPTJForCausalLM.from_pretrained('kakaobrain/kogpt', cache_dir='./my_dir', revision='KoGPT6B-ryan1.5b', torch_dtype=torch.float16)
@@ -11,14 +12,18 @@ model = GPTJForCausalLM.from_pretrained('kakaobrain/kogpt', cache_dir='./my_dir'
 ### fp16 모델 로드 후 문장 생성
 
 ```
+import torch
 from transformers import GPTJForCausalLM, AutoTokenizer
 
-model = GPTJForCausalLM.from_pretrained('MrBananaHuman/kogpt_6b_fp16')
+model = GPTJForCausalLM.from_pretrained('MrBananaHuman/kogpt_6b_fp16', torch_dtype=torch.float16)
 model.to('cuda')
 tokenizer = AutoTokenizer.from_pretrained('MrBananaHuman/kogpt_6b_fp16')
 
 input_text = '이순신은'
-input_ids = tokenizer.encode(input_text, 
+input_ids = tokenizer(input_text, return_tensors='pt').to('cuda')
+
+output = model.generate(**input_ids)
+print(tokenizer.decode(output[0]))
 
 ```