add more samples

app.py

@@ -62,7 +62,7 @@ def get_prediction(model, padded_sequences, img_x, device='cpu'):
         outputs = model(padded_sequences, img_x)
         _, predicted = torch.max(outputs, 1)  
     
-    return malware_classes[predicted]  # Use .item() to get scalar value
+    return malware_classes[predicted]  #
 
 # Define the prediction function for Gradio
 def predict_malware(sha256_hash):
@@ -76,7 +76,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)  # Add batch dimension
+    img_x = img_x.unsqueeze(0)  # type: ignore
 
     # Load the model
     model_path = "model_dump/model_malware_lstm (1).pkl"