Update README.md

README.md

@@ -30,8 +30,10 @@ from transformers import T5ForConditionalGeneration, T5Tokenizer
 model = T5ForConditionalGeneration.from_pretrained("j5ng/et5-typos-corrector")
 tokenizer = T5Tokenizer.from_pretrained("j5ng/et5-typos-corrector")
 
-model = model.to('mps:0') # for mac m1
-# model = model.to('cuda:0') # for nvidia cuda
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+# device = "mps:0" if torch.cuda.is_available() else "cpu" # for mac m1
+
+model = model.to(device) 
 
 # 예시 입력 문장
 input_text = "아늬 진짜 무ㅓ하냐고"
@@ -39,11 +41,8 @@ input_text = "아늬 진짜 무ㅓ하냐고"
 # 입력 문장 인코딩
 input_encoding = tokenizer("맞춤법을 고쳐주세요: " + input_text, return_tensors="pt")
 
-input_ids = input_encoding.input_ids.to('mps:0')
-attention_mask = input_encoding.attention_mask.to('mps:0')
-
-# input_ids = input_encoding.input_ids.to('cuda:0') # for nvidia cuda
-# attention_mask = input_encoding.attention_mask.to('cuda:0') # for nvidia cuda
+input_ids = input_encoding.input_ids.to(device)
+attention_mask = input_encoding.attention_mask.to(device)
 
 # T5 모델 출력 생성
 output_encoding = model.generate(