Add changes

app.py

@@ -20,12 +20,21 @@ class_names = ['apple_pie','baby_back_ribs','baklava','beef_carpaccio','beef_tar
 effnetb4, effnetb4_transforms = create_effnetb4_model(num_classes=len(class_names))
 
 # Load the saved weights
-effnetb4.load_state_dict(
-    torch.load(
-        f="pretrained_efficientnet_b4_model.pth",
-        map_location=torch.device("cpu") # Load the model to the CPU
-    )
-)
+checkpoint = torch.load("pretrained_efficientnet_b4_model.pth", map_location=torch.device("cpu"))
+state_dict = checkpoint
+
+
+# If the state_dict contains "module." due to DataParallel, we need to remove it
+state_dict = checkpoint
+if "module." in list(state_dict.keys())[0]:
+    from collections import OrderedDict
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        new_state_dict[k.replace("module.", "")] = v
+    state_dict = new_state_dict
+
+# Load the state dict into the model
+effnetb4.load_state_dict(state_dict)
 
 # Prediciton function
 def predict(img) -> Tuple[Dict, float]:
@@ -58,7 +67,7 @@ example_list = [["examples/" + example] for example in os.listdir("examples")]
 demo = gr.Interface(fn=predict, # maps inputs to outputs
                    inputs=gr.Image(type="pil"),
                    outputs=[gr.Label(num_top_classes=5, label='Predictions'),
-                           gr.Number(label="Prediction time (s}")],
+                           gr.Number(label="Prediction time (s)")],
                    examples=example_list,
                    title=title,
                    description=description,

model.py

@@ -8,11 +8,8 @@ def create_effnetb4_model(num_classes:int = 101):
     weights = torchvision.models.EfficientNet_B4_Weights.DEFAULT
     transforms = weights.transforms()
 
-    model = torchvision.models.efficientnet_b4(pretrained=False)
+    model = torchvision.models.efficientnet_b4(weights=weights)
 
     model.classifier[1] = nn.Linear(in_features=1792, out_features=num_classes)
 
-    # Use DataParallel to match the saved state dict
-    model = torch.nn.DataParallel(model)
-
     return model, transforms