# create a fucntions that creates an effnet_b2 model and returns its transformation
import torch
import torchvision

from torch import nn
from torchvision import models

def create_effnet_b2():
  # get the weights
  effnet_b2_weights = models.EfficientNet_B2_Weights.DEFAULT
  # get the transforms
  effnet_b2_transforms = effnet_b2_weights.transforms()
  # get the model
  effnet_b2 = models.efficientnet_b2(weights = effnet_b2_weights)

  # freeze all base layers
  for params in effnet_b2.parameters():
    params.requires_grad = False

  # 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 effnet_b2, effnet_b2_transforms