update

model_utils.py

@@ -7,7 +7,6 @@ from NN_classifier.simple_binary_classifier import Medium_Binary_Network
 from feature_extraction import extract_features
 import pandas as pd
 
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 def load_model(model_dir='models/medium_binary_classifier'):
     model_path = os.path.join(model_dir, 'nn_model.pt')
@@ -29,8 +28,8 @@ def load_model(model_dir='models/medium_binary_classifier'):
     
     input_size = scaler.n_features_in_
     
-    model = Medium_Binary_Network(input_size, hidden_sizes=[256, 192, 128, 64], dropout=0.3).to(DEVICE)
-    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
+    model = Medium_Binary_Network(input_size, hidden_sizes=[256, 192, 128, 64], dropout=0.3)
+    model.load_state_dict(torch.load(model_path))
     model.eval()
     
     if imputer is not None:
@@ -76,7 +75,7 @@ def classify_text(text, model, scaler, label_encoder, imputer=None, scores=None)
     
     features_scaled = scaler.transform(features)
     
-    features_tensor = torch.FloatTensor(features_scaled).to(DEVICE)
+    features_tensor = torch.FloatTensor(features_scaled)
     
     with torch.no_grad():
         outputs = model(features_tensor)