Update app.py

app.py

@@ -4,6 +4,8 @@ from transformers import BertTokenizerFast
 
 tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
 model = BertForQuestionAnswering.from_pretrained("bert-base-uncased")
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f'Working on {device}')
 
 def get_prediction(context, question):
   inputs = tokenizer.encode_plus(question, context, return_tensors='pt').to(device)