test

app.py

@@ -14,9 +14,6 @@ from PIL import Image
 from src.utils.get_features import get_img_api
 import joblib 
 
-
-device = torch.device('cpu')
-
 # Path to the dataset
 data_path = 'src/data/subset_dataset.csv'
 device = torch.device('cpu')
@@ -31,13 +28,14 @@ simple_transform = transforms.Compose([
 # Load the model
 def load_model(model_path, device='cpu'):
     """Loads the model from a joblib file and moves it to the specified device."""
-    # Use torch.load with map_location to ensure CPU compatibility
-    with open(model_path, 'rb') as f:
-        model = torch.load(f, map_location=device)
+    # Load the model using joblib
+    model = joblib.load(model_path)
     
-    # If the model is a PyTorch module, move it to the specified device and set to eval mode
+    # If the model is a PyTorch module, move it to the specified device
     if isinstance(model, torch.nn.Module):
+        # Move model to CPU and handle any CUDA tensors
         model = model.to(device)
+        # Set to evaluation mode
         model.eval()
     return model
 
@@ -47,7 +45,8 @@ def get_prediction(model, padded_sequences, img_x, device='cpu'):
                        "Banking Trojan", "Snake Keylogger", "Spyware"]
 
     # Move inputs to the device
-    padded_sequences, img_x = padded_sequences.to(device), img_x.to(device)
+    padded_sequences = padded_sequences.to(device)
+    img_x = img_x.to(device)
     
     # Perform inference
     with torch.no_grad():  # Disable gradient calculation for inference
@@ -68,7 +67,7 @@ def predict_malware(sha256_hash):
     # Load the dataset
     dataset = CombinedDataset(api_call_list, image_path, transforms=simple_transform, sequence_length=config.configuration["sequence_length"])
     padded_sequences, img_x = next(iter(dataset))
-    img_x = img_x.unsqueeze(0) #type: ignore
+    img_x = img_x.unsqueeze(0)  # Add batch dimension
 
     # Load the model
     model_path = "model_dump/model_malware_lstm (1).pkl"