Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import json
 import re
@@ -13,7 +14,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 model.eval()
 if torch.cuda.is_available():
-    model.cuda()
+    model.to("cuda:0")
 
 validation_results = json.load(open('validation_results.json'))
 scores, thresholds, precisions, recalls = validation_results['scores'], validation_results['thresholds'], validation_results['precisions'], validation_results['recalls']
@@ -72,18 +73,24 @@ Authors: {authors}
 Abstract: {abstract}"""
     return text
 
-@torch.no_grad()
+@torch.no_grad
+@spaces.GPU
 def model_inference(title, authors, abstract):
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    if device != model.device:
+        model.to(device)
     text = fill_template(title, authors, abstract)
     text = f'[CLS] {text} [SEP]'
     print (text)
+    print (device)
     inputs = tokenizer([text], return_tensors="pt", truncation=True, max_length=max_length)
     if torch.cuda.is_available():
-        inputs = {key: value.cuda() for key, value in inputs.items()}
+        inputs = {key: value.to(device) for key, value in inputs.items()}
     outputs = model(**inputs)
     logits = outputs.logits
     probs = logits.softmax(dim=-1).view(-1)
     score = probs[1].item()
+    print (score)
     return score
 
 def predict(title, authors, abstract):