model.py

import torchvision
import torch

def create_model(
    num_classes: int=4,
    seed: int=42):


  weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
  transform = weights.transforms()
  model = torchvision.models.efficientnet_b2(weights=weights).to("cpu")

  for params in model.parameters():
    params.requires_grad = False

  model.classifier = torch.nn.Sequential(
      torch.nn.Dropout(p=0.2,inplace=True),
      torch.nn.Linear(1408,num_classes)
  )

  return model, transform