Update app.py

app.py

@@ -6,6 +6,7 @@ import torch.nn.functional as F
 import torch.nn as nn
 import re
 model_path = r'ssocean/NAIP'  
+
 model = AutoModelForSequenceClassification.from_pretrained(model_path, num_labels=1, load_in_8bit=True)
 
 tokenizer = AutoTokenizer.from_pretrained(model_path)
@@ -13,10 +14,12 @@ tokenizer = AutoTokenizer.from_pretrained(model_path)
 
 @spaces.GPU
 def predict(title, abstract):
+    torch.cuda.set_device(torch.device('cuda'))
+    
     model.eval()
     text = f'''Given a certain paper, Title: {title}\n Abstract: {abstract}. \n Predict its normalized academic impact (between 0 and 1):'''
     inputs = tokenizer(text, return_tensors="pt")
-    inputs = inputs.to("cuda")
+    inputs = inputs.to(torch.device('cuda'))
     with torch.no_grad():
         outputs = model(**inputs)
     probability = torch.sigmoid(outputs.logits).item()