code: def __init__(self, args):
      super().__init__()
  
      self.num_queries = args.num_queries
      self.backbone = build_unit_convnet_backbone(args)
      self.transformer = UniTTransformer(args)
      encoder_hidden_dim = self.transformer.d_model_enc
      decoder_hidden_dim = self.transformer.d_model_dec
  
      self.query_embeds = nn.ModuleDict()
      for task_type in self.num_queries:
          task_dict = nn.ModuleDict()
          for dataset in self.num_queries[task_type]:
              task_dict[dataset] = nn.Embedding(
                  self.num_queries[task_type][dataset],
                  decoder_hidden_dim,
              )
          self.query_embeds[task_type] = task_dict
  
      self.input_proj = nn.Conv2d(
          self.backbone.num_channels, encoder_hidden_dim, kernel_size=1
      )
  
imports: from typing import Optional
  import torch
  import torch.nn.functional
  from mmf.models.unit.backbone import build_unit_convnet_backbone
  from mmf.models.unit.matcher import HungarianMatcher
  from mmf.models.unit.misc import NestedTensor
  from mmf.models.unit.transformer import UniTTransformer
  from mmf.utils import box_ops
  from mmf.utils.distributed import get_world_size, is_dist_initialized
  from torch import nn, Tensor
label_desc: Use image node_id 1 to create a new UniTBaseModel object from inherited base classes: nn with example: obj = UniTBaseModel(args)