model.py

import torch
from torchvision import models
from torch import nn
def create_model(num_classes:int=3):
  Eff_Net_Model=models.efficientnet_b2(weights="DEFAULT")
  Eff_Net_Weights=models.EfficientNet_B2_Weights.DEFAULT
  Eff_Net_transform=Eff_Net_Weights.transforms()
  for param in Eff_Net_Model.features.parameters():
    param.requires_grad=False
  Eff_Net_Model.classifier=nn.Sequential(
    nn.Dropout(p=0.3,inplace=True),
    nn.Linear(in_features=1408, out_features=num_classes, bias=True)
  )
  return Eff_Net_Model,Eff_Net_transform