Update app.py

app.py

@@ -13,18 +13,14 @@ 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
-#     )
-# )
-
-# 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")
+effnetb2.load_state_dict(
+    torch.load(
+        f="effnetb2_feature_extractor_food101_mini.pth",
+        map_location=torch.device("cpu") # load the model to the CPU
+    )
+)
+
+
 
 ### 3. Predict function ###
 def predict(img)->Tuple[Dict,float]: