fixed the weight issue

app.py

@@ -146,11 +146,15 @@ model = ResNet50(num_classes=1000)
 try:
     checkpoint = torch.load("best_model.pt", map_location=device)
     if 'model_state_dict' in checkpoint:
-        model.load_state_dict(checkpoint['model_state_dict'])
+        state_dict = checkpoint['model_state_dict']
+        # Remove _orig_mod. prefix if present (from torch.compile)
+        state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
+        model.load_state_dict(state_dict)
         print(f"Model loaded successfully! Top-1 accuracy: {checkpoint.get('top1_accuracy', 'N/A'):.2f}%")
         print(f"Top-5 accuracy: {checkpoint.get('top5_accuracy', 'N/A'):.2f}%")
     else:
-        model.load_state_dict(checkpoint)
+        state_dict = {k.replace('_orig_mod.', ''): v for k, v in checkpoint.items()}
+        model.load_state_dict(state_dict)
         print("Model loaded successfully!")
 except Exception as e:
     print(f"Warning: Could not load model weights: {e}")