Update test.py

test.py

@@ -13,6 +13,9 @@ def get_model_and_tokenizer(model_name):
 default_model_name = "cahya/bert-base-indonesian-522M"
 tokenizer, model = load_model(default_model_name)
 
+# Move model to GPU
+model = model.to('cuda')
+
 # Prediction function
 def predict_hoax(title, content):
     if tokenizer is None or model is None:
@@ -23,6 +26,8 @@ def predict_hoax(title, content):
 
     text = f"{title} [SEP] {content}"
     inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=256)
+    inputs = {key: value.to('cuda') for key, value in inputs.items()}  # Move inputs to GPU
+
     with torch.no_grad():
         outputs = model(**inputs)
     probs = softmax(outputs.logits, dim=1)
@@ -36,6 +41,8 @@ def predict_proba_for_lime(texts):
     results = []
     for text in texts:
         inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=256)
+        inputs = {key: value.to('cuda') for key, value in inputs.items()}  # Move inputs to GPU
+
         with torch.no_grad():
             outputs = model(**inputs)
         probs = softmax(outputs.logits, dim=1).detach().cpu().numpy()