Update app to load from checkpoint

app.py

@@ -9,9 +9,19 @@ from network import create_model  # Import our model architecture
 # Load model from local checkpoint
 def load_model():
     try:
-        model = create_model(num_classes=1000)  # Match your training classes
+        model = create_model(num_classes=1000)
         checkpoint = torch.load('model/model_best.pth', map_location='cpu')
-        model.load_state_dict(checkpoint['model_state_dict'])
+        
+        # Handle DataParallel state dict
+        state_dict = checkpoint['model_state_dict']
+        # Remove 'module.' prefix if it exists
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            name = k.replace('module.', '')  # Remove 'module.' prefix
+            new_state_dict[name] = v
+        
+        # Load the modified state dict
+        model.load_state_dict(new_state_dict)
         model.eval()
         return model
     except Exception as e: