update 3

app.py

@@ -56,12 +56,18 @@ checkpoint_url = f"https://huggingface.co/Maverick98/EcommerceClassifier/resolve
 checkpoint = torch.hub.load_state_dict_from_url(checkpoint_url, map_location=torch.device('cpu'))
 
 # Strip the "module." prefix from the keys in the state_dict if they exist
+# Clean up the state dictionary
+state_dict = checkpoint.get('model_state_dict', checkpoint)
 new_state_dict = {}
-for k, v in checkpoint.items():
+for k, v in state_dict.items():
     if k.startswith("module."):
-        new_state_dict[k[7:]] = v  # Remove "module." prefix
+        new_key = k[7:]  # Remove "module." prefix
     else:
-        new_state_dict[k] = v
+        new_key = k
+
+    # Check if the new_key exists in the model's state_dict, only add if it does
+    if new_key in model.state_dict():
+        new_state_dict[new_key] = v
 
 model.load_state_dict(new_state_dict)