class_docstr: Pooling-based Vision Transformer
  
  A PyTorch implement of 'Rethinking Spatial Dimensions of Vision Transformers'
      - https://arxiv.org/abs/2103.16302
class_docstr_tok: ["Pooling-based","Vision","Transformer","A","PyTorch","implement","of","'Rethinking","Spatial","Dimensions","of","Vision","Transformers'","-","https",":","\/\/arxiv.org\/abs\/2103.16302"]
code: def forward_head(self, x, pre_logits: bool = False) -> torch.Tensor:
      if self.head_dist is not None:
          assert self.global_pool == "token"
          x, x_dist = x[:, 0], x[:, 1]
          x = self.head_drop(x)
          x_dist = self.head_drop(x)
          if not pre_logits:
              x = self.head(x)
              x_dist = self.head_dist(x_dist)
          if (
              self.distilled_training
              and self.training
              and not torch.jit.is_scripting()
          ):
              # only return separate classification predictions when training in distilled mode
              return x, x_dist
          else:
              # during standard train / finetune, inference average the classifier predictions
              return (x + x_dist) / 2
      else:
          if self.global_pool == "token":
              x = x[:, 0]
          x = self.head_drop(x)
          if not pre_logits:
              x = self.head(x)
          return x
  
imports: import math
  import re
  from functools import partial
  from typing import Sequence, Tuple
  import torch
  from torch import nn
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import trunc_normal_, to_2tuple, LayerNorm
  from ._builder import build_model_with_cfg
  from ._registry import register_model, generate_default_cfgs
  from .vision_transformer import Block
label_desc: Use image node_id 9 for calling the PoolingVisionTransformer obj's underlying member method code with example usage: obj.forward_head(x, pre_logits) and returns: x, x,  x_dist, unknown