Update app.py

app.py

@@ -6,13 +6,14 @@ import torch.nn.functional as F
 import torch.nn as nn
 import re
 model_path = r'ssocean/NAIP'  
-
+device = 'cuda:0'
 
 global model, tokenizer
 model = AutoModelForSequenceClassification.from_pretrained(
     model_path,
     num_labels=1,
     load_in_8bit=True,
+    device_map={"": torch.device(device)}
 )
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 model.eval()
@@ -21,9 +22,8 @@ model.eval()
 
 @spaces.GPU(duration=60, enable_queue=True)
 def predict(title, abstract):
-    model_device = next(model.parameters()).device
     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 = tokenizer(text, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = model(**inputs)
     probability = torch.sigmoid(outputs.logits).item()