fixed model architecture error

model.py

@@ -1,26 +1,25 @@
+# create a fucntions that creates an effnet_b2 model and returns its transformation
 import torch
 import torchvision
-from torchvision import models
-from torch import nn
 
-# create a script that will be used to hold the base arch- to load in effnet
+from torch import nn
+from torchvision import models
 
-def create_vit_b_16():
+def create_effnet_b2():
   # get the weights
-  vit_weights = models.ViT_B_16_Weights.DEFAULT
+  effnet_b2_weights = models.EfficientNet_B2_Weights.DEFAULT
   # get the transforms
-  vit_transforms = vit_weights.transforms()
+  effnet_b2_transforms = effnet_b2_weights.transforms()
   # get the model
-  vit_b_16 = models.vit_b_16(weights = vit_weights)
+  effnet_b2 = models.efficientnet_b2(weights = effnet_b2_weights)
 
   # freeze all base layers
-  for param in vit_b_16.parameters():
-    param.requires_grad = False
+  for params in effnet_b2.parameters():
+    params.requires_grad = False
 
-  # change the classifier head
-  vit_b_16.heads = nn.Sequential(nn.Linear(in_features = 768,
-                                          out_features = 3,
-                                          bias = True)
-                                          )
+  # reset the classifier head
+  effnet_b2.classifier = nn.Sequential(nn.Dropout(p = 0.3, inplace = True),
+                                       nn.Linear(in_features = 1408, out_features = 3)
+                                       )
 
-  return vit_b_16, vit_transforms
+  return effnet_b2, effnet_b2_transforms