model.py

import torch
import torchvision
from torch import nn

device='cpu'
def create_model(num_classes: int=13):
  weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT
  transforms = weights.transforms()
  model = torchvision.models.efficientnet_b0(weights=weights).to(device)

  for param in model.features.parameters():
   param.requires_grad = False

  model.classifier = torch.nn.Sequential(
      torch.nn.Dropout(p=0.2, inplace=True),
      torch.nn.Linear(in_features=1280, out_features=13, bias=True)
  ).to(device)

  return model, transforms