Update test.py

test.py

@@ -5,6 +5,8 @@ import numpy as np
 from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
 import streamlit as st
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 @st.cache_resource
 def get_model_and_tokenizer(model_name):
     return load_model(model_name)
@@ -14,7 +16,7 @@ default_model_name = "cahya/bert-base-indonesian-522M"
 tokenizer, model = load_model(default_model_name)
 
 # Move model to GPU
-model = model.to('cuda')
+model = model.to(device)
 
 # Prediction function
 def predict_hoax(title, content):
@@ -26,7 +28,7 @@ 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
+    inputs = {key: value.to(device) for key, value in inputs.items()}  # Move inputs to GPU
 
     with torch.no_grad():
         outputs = model(**inputs)
@@ -41,7 +43,7 @@ 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
+        inputs = {key: value.to(device) for key, value in inputs.items()}  # Move inputs to GPU
 
         with torch.no_grad():
             outputs = model(**inputs)