Update app.py

app.py

@@ -5,10 +5,10 @@ import torch
 from nltk.tokenize import word_tokenize
 import timm
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
 #model = BertForQuestionAnswering.from_pretrained("bert-base-uncased")
-model = BertForQuestionAnswering.from_pretrained("CountingMstar/ai-tutor-bert-model")
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = BertForQuestionAnswering.from_pretrained("CountingMstar/ai-tutor-bert-model").to(device)
 
 def get_prediction(context, question):
   inputs = tokenizer.encode_plus(question, context, return_tensors='pt').to(device)