def METHOD_NAME(self, num_classes, global_pool=None):
    self.num_classes = num_classes
    if global_pool is not None:
        self.pooling = SelectAdaptivePool2d(pool_type=global_pool, flatten=True)
    self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()