code: def get_model(
      model_type: ModelType,
      model_name: str,
      num_classes: int,
      pretrained: bool = True,
      force_reload: bool = False,
      custom_load_fn: Callable = None,
      **kwargs,
  ):
      """
      :param model_type: (ModelType):
          type of model we're trying to obtain (classification or segmentation)
      :param model_name: (string):
          name of the model. By convention (for classification models) lowercase names represent pretrained model variants while Uppercase do not.
      :param num_classes: (int):
          number of classes to initialize with (this will replace the last classification layer or set the number of segmented classes)
      :param pretrained: (bool):
          whether to load the default pretrained version of the model
          NOTE! NOTE! For classification, the lowercase model names are the pretrained variants while the Uppercase model names are not.
          The only exception applies to torch.hub models (all efficientnet, mixnet, mobilenetv3, mnasnet, spnasnet variants) where a single
          lower-case string can be used for vanilla and pretrained versions. Otherwise, it is IN ERROR to specify an Uppercase model name variant
          with pretrained=True but one can specify a lowercase model variant with pretrained=False
          (default: True)
      :param force_reload: (bool):
          Whether to force reloading the list of models from torch.hub. By default, a cache file is used if it is found locally and that can prevent
          new or updated models from being found.
      :param custom_load_fn: (Callable):
          A custom callable function to use for loading models (typically used to load cutting-edge or custom models that are not in the publicly available list)
  
      :return: model
      """
  
      if (
          model_name not in get_supported_models(model_type)
          and not model_name.startswith("TEST")
          and custom_load_fn is None
      ):
          raise ValueError(
              f"The supplied model name: {model_name} was not found in the list of acceptable model names."
              " Use get_supported_models() to obtain a list of supported models or supply a custom_load_fn"
          )
  
      print(
          "INFO: Loading Model:   --   "
          + model_name
          + "  with number of classes: "
          + str(num_classes)
      )
  
      if model_type == ModelType.CLASSIFICATION:
          torch_hub_names = torch.hub.list(
              rwightman_repo, force_reload=force_reload
          )
          if model_name in torch_hub_names:
              model = torch.hub.load(
                  rwightman_repo,
                  model_name,
                  pretrained=pretrained,
                  num_classes=num_classes,
              )
          elif custom_load_fn is not None:
              model = custom_load_fn(
                  model_name, pretrained, num_classes, **kwargs
              )
          else:
              # 1. Load model (pretrained or vanilla)
              import ssl
  
              ssl._create_default_https_context = (
                  ssl._create_unverified_context
              )
              fc_name = get_fc_names(
                  model_name=model_name, model_type=model_type
              )[-1:][
                  0
              ]  # we're only interested in the last layer name
              new_fc = None  # Custom layer to replace with (if none, then it will be handled generically)
              if model_name in torch_models.__dict__:
                  print(
                      "INFO: Loading torchvision model: {}\t Pretrained: {}".format(
                          model_name, pretrained
                      )
                  )
                  model = torch_models.__dict__[model_name](
                      pretrained=pretrained
                  )  # find a model included in the torchvision package
              else:
                  net_list = [
                      "fbresnet",
                      "inception",
                      "mobilenet",
                      "nasnet",
                      "polynet",
                      "resnext",
                      "se_resnet",
                      "senet",
                      "shufflenet",
                      "xception",
                  ]
                  if pretrained:
                      print(
                          "INFO: Loading a pretrained model: {}".format(
                              model_name
                          )
                      )
                      if "dpn" in model_name:
                          model = classification.__dict__[model_name](
                              pretrained=True
                          )  # find a model included in the pywick classification package
                      elif any(
                          net_name in model_name
                          for net_name in net_list
                      ):
                          model = classification.__dict__[model_name](
                              pretrained="imagenet"
                          )
                  else:
                      print(
                          "INFO: Loading a vanilla model: {}".format(
                              model_name
                          )
                      )
                      model = classification.__dict__[model_name](
                          pretrained=None
                      )  # pretrained must be set to None for the extra models... go figure
  
              # 2. Create custom FC layers for non-standardized models
              if "squeezenet" in model_name:
                  final_conv = nn.Conv2d(
                      512, num_classes, kernel_size=1
                  )
                  new_fc = nn.Sequential(
                      nn.Dropout(p=0.5),
                      final_conv,
                      nn.ReLU(inplace=True),
                      nn.AvgPool2d(13, stride=1),
                  )
                  model.num_classes = num_classes
              elif "vgg" in model_name:
                  new_fc = nn.Sequential(
                      nn.Linear(512 * 7 * 7, 4096),
                      nn.ReLU(True),
                      nn.Dropout(),
                      nn.Linear(4096, 4096),
                      nn.ReLU(True),
                      nn.Dropout(),
                      nn.Linear(4096, num_classes),
                  )
              elif (
                  "inception3" in model_name.lower()
                  or "inception_v3" in model_name.lower()
              ):
                  # Replace the extra aux_logits FC layer if aux_logits are enabled
                  if getattr(model, "aux_logits", False):
                      model.AuxLogits = InceptionAux(768, num_classes)
              elif "dpn" in model_name.lower():
                  old_fc = getattr(model, fc_name)
                  new_fc = nn.Conv2d(
                      old_fc.in_channels,
                      num_classes,
                      kernel_size=1,
                      bias=True,
                  )
  
              # 3. For standard FC layers (nn.Linear) perform a reflection lookup and generate a new FC
              if new_fc is None:
                  old_fc = getattr(model, fc_name)
                  new_fc = nn.Linear(old_fc.in_features, num_classes)
  
              # 4. perform replacement of the last FC / Linear layer with a new one
              setattr(model, fc_name, new_fc)
  
          return model
  
      elif model_type == ModelType.SEGMENTATION:
          """
          Additional Segmentation Option Parameters
          -----------------------------------------
  
          BiSeNet
              - :param backbone: (str, default: 'resnet18') The type of backbone to use (one of `{'resnet18'}`)
              - :param aux: (bool, default: False) Whether to output auxiliary loss (typically an FC loss to help with multi-class segmentation)
  
          DANet_ResnetXXX, DUNet_ResnetXXX, EncNet, OCNet_XXX_XXX, PSANet_XXX
              - :param aux: (bool, default: False) Whether to output auxiliary loss (typically an FC loss to help with multi-class segmentation)
              - :param backbone: (str, default: 'resnet101') The type of backbone to use (one of `{'resnet50', 'resnet101', 'resnet152'}`)
              - :param norm_layer (Pytorch nn.Module, default: nn.BatchNorm2d) The normalization layer to use. Typically it is not necessary to change this parameter unless you know what you're doing.
  
          DenseASPP_XXX
              - :param aux: (bool, default: False) Whether to output auxiliary loss (typically an FC loss to help with multi-class segmentation)
              - :param backbone: (str, default: 'densenet161') The type of backbone to use (one of `{'densenet121', 'densenet161', 'densenet169', 'densenet201'}`)
              - :param dilate_scale (int, default: 8) The size of the dilation to use (one of `{8, 16}`)
              - :param norm_layer (Pytorch nn.Module, default: nn.BatchNorm2d) The normalization layer to use. Typically it is not necessary to change this parameter unless you know what you're doing.
  
          DRNSeg
              - :param model_name: (str - required) The type of backbone to use. One of `{'DRN_C_42', 'DRN_C_58', 'DRN_D_38', 'DRN_D_54', 'DRN_D_105'}`
  
          EncNet_ResnetXXX
              - :param aux: (bool, default: False) Whether to output auxiliary loss (typically an FC loss to help with multi-class segmentation)
              - :param backbone: (str, default: 'resnet101') The type of backbone to use (one of `{'resnet50', 'resnet101', 'resnet152'}`)
              - :param norm_layer (Pytorch nn.Module, default: nn.BatchNorm2d) The normalization layer to use. Typically it is not necessary to change this parameter unless you know what you're doing.
              - :param se_loss (bool, default: True) Whether to compute se_loss
              - :param lateral (bool, default: False)
  
          frrn
              - :param model_type: (str - required) The type of model to use. One of `{'A', 'B'}`
  
          GCN, GCN_DENSENET, GCN_NASNET, GCN_PSP, GCN_RESNEXT
              - :param k: (int - optional) The size of global kernel
  
          GCN_PSP, GCN_RESNEXT, Unet_stack
              - :param input_size: (int - required) The size of output image (will be square)
  
          LinkCeption, 'LinkDenseNet121', 'LinkDenseNet161', 'LinkInceptionResNet', 'LinkNet18', 'LinkNet34', 'LinkNet50', 'LinkNet101', 'LinkNet152', 'LinkNeXt', 'CoarseLinkNet50'
              - :param num_channels: (int, default: 3) Number of channels in the image (e.g. 3 = RGB)
              - :param is_deconv: (bool, default: False)
              - :param decoder_kernel_size: (int, default: 3) Size of the decoder kernel
  
          PSPNet
              - :param backend: (str, default: densenet121) The type of extractor to use. One of `{'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'densenet121'}`
  
          RefineNet4Cascade, RefineNet4CascadePoolingImproved
              - :param input_shape: (tuple(int, int), default: (1, 512) - required!) Tuple representing input shape (num_channels, dim)
              - :param freeze_resnet: (bool, default: False) - whether to freeze the underlying resnet
          """
  
          model_exists = False
          for m_name in get_supported_models(model_type):
              if model_name in m_name:
                  model_exists = True
                  break
          if model_exists:
              # Print warnings and helpful messages for nets that require additional configuration
              if model_name in [
                  "GCN_PSP",
                  "GCN_RESNEXT",
                  "RefineNet4Cascade",
                  "RefineNet4CascadePoolingImproved",
                  "Unet_stack",
              ]:
                  print(
                      "WARN: Did you remember to set the input_size parameter: (int) ?"
                  )
              elif model_name in [
                  "RefineNet4Cascade",
                  "RefineNet4CascadePoolingImproved",
              ]:
                  print(
                      "WARN: Did you remember to set the input_shape parameter: tuple(int, int)?"
                  )
  
              # logic to switch between different constructors
              if (
                  model_name
                  in [
                      "FusionNet",
                      "Enet",
                      "frrn",
                      "Tiramisu57",
                      "Tiramisu67",
                      "Tiramisu101",
                  ]
                  or model_name.startswith("UNet")
                  and pretrained
              ):  # FusionNet
                  print(
                      "WARN: FusionNet, Enet, FRRN, Tiramisu, UNetXXX do not have a pretrained model! Empty model as been created instead."
                  )
  
              net = segmentation.__dict__[model_name](
                  num_classes=num_classes,
                  pretrained=pretrained,
                  **kwargs,
              )
  
          else:
              raise Exception(
                  "Combination of type: {} and model_name: {} is not valid".format(
                      model_type, model_name
                  )
              )
  
      return net
  
imports: from functools import partial
  from typing import Callable
  from torchvision.models.resnet import Bottleneck
  from .None import classification
  from .segmentation import *
  from .None import segmentation
  from enum import Enum
  from torchvision import models
  from torchvision.models.inception import InceptionAux
  import torch
  import torch.nn
  import os
  import errno
label_desc: Use image node_id 2 for calling a global function with example usage: get_model(model_type, model_name, num_classes, pretrained, force_reload, custom_load_fn) and returns: net, model