import torch
import torchvision
from torch import nn

from torchvision.models._api import WeightsEnum
from torch.hub import load_state_dict_from_url
def get_state_dict(self, *args, **kwargs):
    kwargs.pop("check_hash")
    return load_state_dict_from_url(self.url, *args, **kwargs)
WeightsEnum.get_state_dict = get_state_dict

def create_effnetb2_model(num_classes : int = 3,
                          seed : int = 42):

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


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

  torch.manual_seed(seed)
  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