final?

app.py

@@ -12,6 +12,7 @@ import pickle
 import torch
 from PIL import Image
 from src.utils.get_features import get_img_api
+import joblib 
 
 # Path to the dataset
 data_path = 'src/data/subset_dataset.csv'
@@ -26,10 +27,12 @@ simple_transform = transforms.Compose([
 
 # Load the model
 def load_model(model_path, device='cpu'):
-    """Loads the model from a pickle file and moves it to the specified device."""
-    with open(model_path, 'rb') as f:
-        model = pickle.load(f)
-    return model.to(device)
+    """Loads the model from a joblib file and moves it to the specified device."""
+    model = joblib.load(model_path)
+    # If the model contains PyTorch tensors, move them to the specified device
+    if isinstance(model, torch.nn.Module):
+        model = model.to(device)
+    return model
 
 # Get prediction
 def get_prediction(model, padded_sequences, img_x, device='cpu'):