code: def __init__(
      self,
      dim: int,
      dim_out: int,
      depth: int,
      downsample: bool = True,
      num_heads: int = 8,
      sr_ratio: int = 1,
      linear_attn: bool = False,
      mlp_ratio: float = 4.0,
      qkv_bias: bool = True,
      proj_drop: float = 0.0,
      attn_drop: float = 0.0,
      drop_path: Union[List[float], float] = 0.0,
      norm_layer: Callable = LayerNorm,
  ):
      super().__init__()
      self.grad_checkpointing = False
  
      if downsample:
          self.downsample = OverlapPatchEmbed(
              patch_size=3,
              stride=2,
              in_chans=dim,
              embed_dim=dim_out,
          )
      else:
          assert dim == dim_out
          self.downsample = None
  
      self.blocks = nn.ModuleList(
          [
              Block(
                  dim=dim_out,
                  num_heads=num_heads,
                  sr_ratio=sr_ratio,
                  linear_attn=linear_attn,
                  mlp_ratio=mlp_ratio,
                  qkv_bias=qkv_bias,
                  proj_drop=proj_drop,
                  attn_drop=attn_drop,
                  drop_path=drop_path[i]
                  if isinstance(drop_path, list)
                  else drop_path,
                  norm_layer=norm_layer,
              )
              for i in range(depth)
          ]
      )
  
      self.norm = norm_layer(dim_out)
  
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 1 to create a new PyramidVisionTransformerStage object from inherited base classes: nn with example: obj = PyramidVisionTransformerStage(dim, dim_out, depth, downsample, num_heads, sr_ratio, linear_attn, mlp_ratio, qkv_bias, proj_drop, attn_drop, drop_path, norm_layer)