docstr: assuming input size 7x7 or 14x14
docstr_tok: ["assuming","input","size","7x7","or","14x14"]
code: def __init__(self, C, size, num_classes, bn_affine=False):
      """assuming input size 7x7 or 14x14"""
      super(DistillHeadImagenet, self).__init__()
      self.features = nn.Sequential(
          nn.ReLU(),
          nn.AvgPool2d(
              size, stride=2, padding=0, count_include_pad=False
          ),  # image size = 2 x 2 / 6 x 6
          nn.Conv2d(C, 128, 1, bias=False),
          nn.BatchNorm2d(128, affine=bn_affine),
          nn.ReLU(),
          nn.Conv2d(128, 768, 2, bias=False),
          nn.BatchNorm2d(768, affine=bn_affine),
          nn.ReLU(),
      )
      self.classifier = nn.Linear(768, num_classes)
      self.gap = nn.AdaptiveAvgPool2d(1)
  
file_path: nni/examples/nas/legacy/cdarts/aux_head.py
imports: import torch.nn
label_desc: Use image node_id 1 to create a new DistillHeadImagenet object from inherited base classes: nn with example: obj = DistillHeadImagenet(C, size, num_classes, bn_affine)