test

app.py

@@ -12,7 +12,17 @@ import pickle
 import torch
 from PIL import Image
 from src.utils.get_features import get_img_api
-import joblib 
+import joblib
+import io
+
+# Custom unpickler to handle device mapping
+class CPU_Unpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == "torch.storage" and name == "_load_from_bytes":
+            def _load_from_bytes(b):
+                return torch.load(io.BytesIO(b), map_location=torch.device('cpu'))
+            return _load_from_bytes
+        return super().find_class(module, name)
 
 # Path to the dataset
 data_path = 'src/data/subset_dataset.csv'
@@ -27,16 +37,15 @@ 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."""
-    # Load the model using joblib
-    model = joblib.load(model_path)
+    """Loads the model from a joblib file and ensures it runs on the specified device."""
+    # Load the model using joblib with custom unpickler
+    with open(model_path, 'rb') as f:
+        model = CPU_Unpickler(f).load()
     
     # 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()
+        model.eval()  # Set to evaluation mode
     return model
 
 # Get prediction