code: def forward(self, x):
      # x is either B, C, H, W (if downsample) or B, H, W, C if not
      if self.downsample is not None:
          # input to downsample is B, C, H, W
          x = self.downsample(x)  # output B, H, W, C
      B, H, W, C = x.shape
      feat_size = (H, W)
      x = x.reshape(B, -1, C)
      for blk in self.blocks:
          if self.grad_checkpointing and not torch.jit.is_scripting():
              x = checkpoint.checkpoint(blk, x, feat_size)
          else:
              x = blk(x, feat_size)
      x = self.norm(x)
      x = (
          x.reshape(B, feat_size[0], feat_size[1], -1)
          .permute(0, 3, 1, 2)
          .contiguous()
      )
      return x
  
imports: import math
  from typing import Tuple, List, Callable, Union
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils.checkpoint
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import DropPath, to_2tuple, to_ntuple, trunc_normal_, LayerNorm, use_fused_attn
  from ._builder import build_model_with_cfg
  from ._registry import register_model, generate_default_cfgs
label_desc: Use image node_id 2 for calling the PyramidVisionTransformerStage obj's underlying member method code with example usage: obj.forward(x) and returns: x