Update app.py

app.py

@@ -2,13 +2,14 @@ import streamlit as st
 import torch
 from transformers import AutoModelForQuestionAnswering, AutoTokenizer
 
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
 @st.cache(allow_output_mutation=True)
 def get_model():
     # Load fine-tuned MRC model by HuggingFace Model Hub
     HUGGINGFACE_MODEL_PATH = "bespin-global/klue-bert-base-aihub-mrc"
     tokenizer = AutoTokenizer.from_pretrained(HUGGINGFACE_MODEL_PATH)
-    model = AutoModelForQuestionAnswering.from_pretrained(HUGGINGFACE_MODEL_PATH)
+    model = AutoModelForQuestionAnswering.from_pretrained(HUGGINGFACE_MODEL_PATH).to(device)
 
     return tokenizer, model
 
@@ -62,7 +63,7 @@ if st.button("Submit", key='question'):
                               return_token_type_ids=False,
                               return_offsets_mapping=True
                               )
-            encodings = {key: torch.tensor([val]) for key, val in encodings.items()}             
+            encodings = {key: torch.tensor([val]).to(device) for key, val in encodings.items()}             
 
             # Predict
             pred = model(encodings["input_ids"], attention_mask=encodings["attention_mask"])