utils/load_model.py

import torch
import torchvision

def load_model(architecture:str="resnet_50",
               pretrained:bool=True) -> torch.nn.Module:
  """
    Load a pre-trained/untrained model based on specified architecture and set it to evaluation mode.
    Args:
        architecture (str): Name of the model architecture (e.g. 'resnet50', 'vgg16').
    Returns:
        torch.nn.Module: Pre-trained model.
  """
  model_dict = {
      "resnet50": (torchvision.models.resnet50, torchvision.models.ResNet50_Weights.DEFAULT),
      "vgg16": (torchvision.models.vgg16, torchvision.models.VGG16_Weights.DEFAULT),
      "mobilenet_v2": (torchvision.models.mobilenet_v2, torchvision.models.MobileNet_V2_Weights.DEFAULT),
      "efficientnet_b0": (torchvision.models.efficientnet_b0, torchvision.models.EfficientNet_B0_Weights.DEFAULT)
      #Feel free to add more architectures here
  }

  try:

    # Retrieve the model function and weights from model_dict
    model_fn, weights = model_dict[architecture]

    # Use weights if pretrained, otherwise use None
    weights = weights if pretrained else None

    # Load the model in eval mode
    model = model_fn(weights=weights)
    model.eval()
    return model

  except KeyError:
    raise ValueError(f"Unsupported architecture {architecture}. Available options{list(model_dict.keys())}")

  except Exception as e:
    raise RuntimeError(f"An error occurred while loading the model {e}")