model.py

import torch
import torchvision
def create_effnet_b2_model(num_classes=101):
  """
  Args: num_classes is total number of classes

  returns: model and its corresponding model_specific transform
  """
  weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
  transform = torchvision.models.EfficientNet_B2_Weights.DEFAULT.transforms()
  model = torchvision.models.efficientnet_b2(weights=weights)

  # freeze the parameters from training
  for param in model.parameters():
    param.requires_grad = False

  # modifying the classifier layer
  model.classifier = torch.nn.Sequential(
      torch.nn.Dropout(p=0.3,inplace=True),
      torch.nn.Linear(in_features=1408,out_features=num_classes)
  )

  return model,transform