use cuda:0 as torch device

interfaces/cap.py

@@ -85,14 +85,15 @@ def build_huggingface_path(language: str, domain: str):
 
 @spaces.GPU(duration=5)
 def predict(text, model_id, tokenizer_id):
-    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", offload_folder="offload", token=HF_TOKEN)
+    device = torch.device("cuda:0")
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", offload_folder="offload", token=HF_TOKEN).to(device)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
     inputs = tokenizer(text,
                        max_length=256,
                        truncation=True,
                        padding="do_not_pad",
-                       return_tensors="pt")
+                       return_tensors="pt").to(device)
     model.eval()
 
     with torch.no_grad():

interfaces/cap_minor_media.py

@@ -59,6 +59,8 @@ def build_huggingface_path(language: str, domain: str):
 
 @spaces.GPU(duration=5)
 def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
+    device = torch.device("cuda:0")
+
     # Load major and minor models + tokenizer
     major_model = AutoModelForSequenceClassification.from_pretrained(
         major_model_id,
@@ -66,7 +68,7 @@ def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
         device_map="auto",
         offload_folder="offload",
         token=HF_TOKEN
-    )
+    ).to(device)
 
     minor_model = AutoModelForSequenceClassification.from_pretrained(
         minor_model_id,
@@ -74,12 +76,12 @@ def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
         device_map="auto",
         offload_folder="offload",
         token=HF_TOKEN
-    )
+    ).to(device)
 
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
     # Tokenize input
-    inputs = tokenizer(text, max_length=256, truncation=True, padding="do_not_pad", return_tensors="pt")
+    inputs = tokenizer(text, max_length=256, truncation=True, padding="do_not_pad", return_tensors="pt").to(device)
 
     # Predict major topic
     major_model.eval()