Update app.py

app.py

@@ -1,9 +1,11 @@
 import gradio as gr
+import spaces
 from transformers import pipeline
 
-token_skill_classifier = pipeline(model="jjzha/jobbert_skill_extraction", aggregation_strategy="first", device="cuda")
-token_knowledge_classifier = pipeline(model="jjzha/jobbert_knowledge_extraction", aggregation_strategy="first", device="cuda")
-
+token_skill_classifier = pipeline(model="jjzha/jobbert_skill_extraction", aggregation_strategy="first")
+token_knowledge_classifier = pipeline(model="jjzha/jobbert_knowledge_extraction", aggregation_strategy="first")
+token_skill_classifier.to("cuda")
+token_knowledge_classifier.to("cuda")
 
 examples = [
         "Knowing Python is a plus",
@@ -27,7 +29,7 @@ def aggregate_span(results):
 
     return new_results
 
-
+@spaces.GPU
 def ner(text):
     output_skills = token_skill_classifier(text)
     for result in output_skills: