fix: use cuda as device

app.py

@@ -1,21 +1,24 @@
-import gradio as gr
+import torch
 import spaces
+import gradio as gr
+
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     AutoModelForSequenceClassification,
 )
-import torch
 
 chat_model_name = "sapienzanlp/Minerva-7B-instruct-v1.0"
-chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, dtype=torch.bfloat16, device_map="auto")
+chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, dtype=torch.bfloat16)
+chat_model.to("cuda")
 chat_tokenizer = AutoTokenizer.from_pretrained(chat_model_name)
 
 moderator_model_name = "saiteki-kai/QA-DeBERTa-v3-large"
-moderator_model = AutoModelForSequenceClassification.from_pretrained(moderator_model_name, device_map="auto")
+moderator_model = AutoModelForSequenceClassification.from_pretrained(moderator_model_name)
+moderator_model.to("cuda")
 moderator_tokenizer = AutoTokenizer.from_pretrained(moderator_model_name)
 
-@spaces.GPU()
+@spaces.GPU
 def generate_responses(model, tokenizer, prompts):
     messages = [[{"role": "user", "content": message}] for message in prompts]
 
@@ -38,7 +41,7 @@ def generate_responses(model, tokenizer, prompts):
 
     return responses
 
-@spaces.GPU()
+@spaces.GPU
 def classify_pairs(model, tokenizer, prompts, responses):
     texts = [
         prompt + "[SEP]" + response for prompt, response in zip(prompts, responses)
@@ -56,7 +59,7 @@ def classify_pairs(model, tokenizer, prompts, responses):
     return unsafety_scores
 
 
-@spaces.GPU(duration=60)
+@spaces.GPU
 def generate(submission: list[dict[str, str]], team_id: str) -> list[dict[str, str | float]]:
     ids = [s["id"] for s in submission]
     prompts = [s["prompt"] for s in submission]