model.py

import torch
import torchvision
from torch import nn

def create_effnetb2_model(num_classes: int=101,
                          seed: int=29):
  # creating pretrained weights
  weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT

  # setting up the transforms
  transforms = weights.transforms()

  # creating the model
  model = torchvision.models.efficientnet_b2(weights=weights)

  # freezing all the base layers
  for param in model.parameters():
    param.requires_grad = False

  # changing the cloassifier head
  torch.manual_seed(seed)
  torch.cuda.manual_seed(seed)
  model.classifier = nn.Sequential(
    nn.Dropout(p=0.3,
               inplace=True),
    nn.Linear(in_features=1408,
              out_features=num_classes)
  )

  return model, transforms