code: def _omnivore_base(
      trunk: nn.Module,
      heads: Optional[Union[nn.Module, nn.ModuleDict]] = None,
      head_dim_in: int = 1024,
      pretrained: bool = True,
      progress: bool = True,
      load_heads: bool = True,
      checkpoint_name: str = "omnivore_swinB",
  ) -> nn.Module:
      """
      Load and initialize the specified Omnivore
      model trunk (and optionally heads).
      Args:
          trunk: nn.Module of the SwinTransformer3D trunk
          heads: Provide the heads module if using a custom
              model. If not provided image/video/rgbd heads are
              added corresponding to the omnivore base model.
          head_dim_in: Only needs to be set if heads = None.
              The dim is used for the default base model heads.
          load_heads: if True, loads the 3 heads, one each for
              image/video/rgbd prediction. If False loads only the
              trunk.
      Returns:
          model: nn.Module of the full Omnivore model
      """
      if load_heads and heads is None:
          # Get heads
          heads = get_all_heads(dim_in=head_dim_in)
  
      if pretrained:
          path = CHECKPOINT_PATHS[checkpoint_name]
  
          # All models are loaded onto CPU by default
          checkpoint = load_state_dict_from_url(
              path, progress=progress, map_location="cpu"
          )
          trunk.fc_cls = nn.Sequential()
          trunk.load_state_dict(checkpoint["trunk"])
  
          if load_heads:
              heads.load_state_dict(checkpoint["heads"])
  
      if load_heads:
          model = OmnivoreModel(trunk=trunk, heads=heads)
      else:
          model = trunk
  
      return model
  
imports: from towhee.models.video_swin_transformer.video_swin_transformer import VideoSwinTransformer
  from typing import Any, Optional, Union
  import torch
  from torch import nn
  from torch.hub import load_state_dict_from_url
label_desc: Use image node_id 5 for calling a global function with example usage: _omnivore_base(trunk, heads, head_dim_in, pretrained, progress, load_heads, checkpoint_name) and returns: model