fixed state_dict

app.py

@@ -18,7 +18,15 @@ vit16, vit16_transforms = create_vit16_model(
     num_classes=101, # could also use len(class_names)
 )
 
-# Load saved weights
+
+state_dict = torch.load("model_food101_20_percent.pth")
+state_dict["heads.0.weight"] = state_dict.pop("heads.weight")
+state_dict["heads.0.bias"] = state_dict.pop("heads.bias")
+# save new state_dict in .pth
+torch.save(state_dict, "model_food101_20_percent.pth")
+
+
+
 vit16.load_state_dict(
     torch.load(
         f="model_food101_20_percent.pth",
@@ -26,6 +34,8 @@ vit16.load_state_dict(
     )
 )
 
+
+
 ### 3. Predict function ###
 
 # Create predict function

model.py

@@ -28,8 +28,7 @@ def create_vit16_model(num_classes:int=101,
 
     # Change classifier head with random seed for reproducibility
     torch.manual_seed(seed)
-    model.classifier = nn.Sequential(
-        nn.Linear(in_features=768, out_features=num_classes).to("cpu"),
-    )
+    model.heads = nn.Sequential(nn.Linear(in_features=768, # keep this the same as original model
+                                          out_features=num_classes)) # update to reflect target number of classes
     
     return model, transforms

model_food101_20_percent.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:26afc8e7a4a879d88f402030c346e073fdc2671ed6928c87305a1e2776f72c5b
+oid sha256:b4357a334ed5737baacaf7a99b0ba491ef88c61580790277acec9ef877cd77c9
 size 343564561