class_docstr: Auxiliary head in 2/3 place of network to let the gradient flow well 
docstr: assuming input size 7x7 or 8x8 
docstr_tok: ["assuming","input","size","7x7","or","8x8"]
code: def __init__(self, input_size, C, n_classes):
      """assuming input size 7x7 or 8x8"""
      assert input_size in [7, 8]
      super().__init__()
      self.net = nn.Sequential(
          nn.ReLU(inplace=True),
          nn.AvgPool2d(
              5,
              stride=input_size - 5,
              padding=0,
              count_include_pad=False,
          ),  # 2x2 out
          nn.Conv2d(C, 128, kernel_size=1, bias=False),
          nn.BatchNorm2d(128),
          nn.ReLU(inplace=True),
          nn.Conv2d(128, 768, kernel_size=2, bias=False),  # 1x1 out
          nn.BatchNorm2d(768),
          nn.ReLU(inplace=True),
      )
      self.linear = nn.Linear(768, n_classes)
  
imports: from collections import OrderedDict
  import torch
  import torch.nn
  import ops
  from nni.retiarii.nn.pytorch import LayerChoice, InputChoice
label_desc: Use image node_id 1 to create a new AuxiliaryHead object from inherited base classes: nn with example: obj = AuxiliaryHead(input_size, C, n_classes)