model.py

import torch
import torchvision


from torch import nn

def effnet_model(num_classes:int=101,
              seed:int=42):
  
  weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
  transforms = weights.transforms()
  model = torchvision.models.efficientnet_b2(weights=weights)


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

  torch.manual_seed(seed)

  model.classifier = nn.Sequential(
                                  nn.Dropout(p=0.3,inplace=True),
                                  nn.Linear(in_features=1408,out_features=101,bias=True)
                                  )
  return model,transforms