Update model.py

model.py

@@ -3,22 +3,20 @@ import torchvision
 
 from torch import nn
 
-def create_transformer_model(num_classes:int=325, # default output classes = 3 (pizza, steak, sushi)
-                          seed:int=42):
-  # 1, 2, 3 Create EffNetB2 pretrained weights, transforms and model
+def create_vit_model(num_classes:int=325,
+                     seed:int=42):
+  # Create ViT_B_16 pretrained weights, transforms and model
   weights = torchvision.models.ViT_B_16_Weights.DEFAULT
   transforms = weights.transforms()
   model = torchvision.models.vit_b_16(weights=weights)
 
-  # 4. Freeze all layers in the base model
+  # Freeze all of the base layers
   for param in model.parameters():
     param.requires_grad = False
 
-  # 5. Change classifier head with random seed for reproducibility
+  # Change classifier head to suit our needs
   torch.manual_seed(seed)
-  model.classifier = nn.Sequential(
-      nn.Dropout(p=0.3, inplace=True),
-      nn.Linear(in_features=768, out_features=num_classes)
-  )
-
-  return model, transforms
+  model.heads = nn.Sequential(nn.Linear(in_features=768, 
+                                        out_features=325))
+  
+  return model, transforms