code: def __init__(
      self,
      in_unit: int = 368,
      out_unit: int = 1000,
      output_size: tuple = (1, 1),
      drop_ratio: float = 0.25,
  ):
      super(RegHead, self).__init__()
      self.pool = nn.AdaptiveAvgPool2d(output_size)
  
      if drop_ratio > 0:
          self.dropout = nn.Dropout(p=drop_ratio)
      else:
          self.dropout = nn.Identity()
  
      self.fc = nn.Linear(in_features=in_unit, out_features=out_unit)
  
imports: from typing import Optional
  import numpy
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 1 to create a new RegHead object from inherited base classes: nn with example: obj = RegHead(in_unit, out_unit, output_size, drop_ratio)