Update app.py

app.py

@@ -13,12 +13,18 @@ class_names=["pizza","steak","sushi"]
 effnetb2,effnetb2_transforms=create_effnetb2_model(num_classes=len(class_names))
 
 # Load save weights
-effnetb2.load_state_dict(
-    torch.load(
-        f="effnetb2_feature_extractor_food101_mini.pth",
-        map_location=torch.device("cpu") # load the model to the CPU
-    )
-)
+# effnetb2.load_state_dict(
+#     torch.load(
+#         f="effnetb2_feature_extractor_food101_mini.pth",
+#         map_location=torch.device("cpu") # load the model to the CPU
+#     )
+# )
+
+# Load the state_dict (no DataParallel prefix)
+effnetb2.load_state_dict(torch.load("effnetb2_feature_extractor_food101_mini.pth", map_location=torch.device("cpu")))
+
+# Wrap with DataParallel if needed
+# effnetb2 = DataParallel(effnetb2).to("device")
 
 ### 3. Predict function ###
 def predict(img)->Tuple[Dict,float]: