class_docstr: WaveViT
  Args:
      in_chans (`int`):
          Number of image channels.
      num_classes (`int`):
          Number of image classes.
      stem_hidden_dim (`int`):
          Stem layer hidden dim.
      embed_dims (`list`):
          Every layer`s embedding dim.
      num_heads (`list`):
          Every layer`s head number.
      mlp_ratios (`list`):
          Every layer`s mlp ration.
      drop_path_rate (`float`):
          Every layer`s drop rate.
      norm_layer (`nn.Module`):
          Default nn.LayerNorm.
      depths (`int`):
          Number of depths.
      sr_ratios (`list`):
          Every block`s sr ration.
      num_stages (`int`):
          Number of stages.
      token_label (`bool`):
          Whether to use the token_label.
class_docstr_tok: ["WaveViT","Args",":","in_chans","(","`","int","`",")",":","Number","of","image","channels",".","num_classes","(","`","int","`",")",":","Number","of","image","classes",".","stem_hidden_dim","(","`","int","`",")",":","Stem","layer","hidden","dim",".","embed_dims","(","`","list","`",")",":","Every","layer","`","s","embedding","dim",".","num_heads","(","`","list","`",")",":","Every","layer","`","s","head","number",".","mlp_ratios","(","`","list","`",")",":","Every","layer","`","s","mlp","ration",".","drop_path_rate","(","`","float","`",")",":","Every","layer","`","s","drop","rate",".","norm_layer","(","`","nn.Module","`",")",":","Default","nn.LayerNorm",".","depths","(","`","int","`",")",":","Number","of","depths",".","sr_ratios","(","`","list","`",")",":","Every","block","`","s","sr","ration",".","num_stages","(","`","int","`",")",":","Number","of","stages",".","token_label","(","`","bool","`",")",":","Whether","to","use","the","token_label","."]
code: def forward(self, x):
      if not self.return_dense:
          x = self.forward_features(x)
          x = self.head(x)
          return x
      else:
          x, h, w = self.forward_embeddings(x)
          # mix token, see token labeling for details.
          if self.mix_token and self.training:
              lam = np.random.beta(self.beta, self.beta)
              patch_h, patch_w = (
                  x.shape[1] // self.pooling_scale,
                  x.shape[2] // self.pooling_scale,
              )
              bbx1, bby1, bbx2, bby2 = rand_bbox(
                  x.size(), lam, scale=self.pooling_scale
              )
              temp_x = x.clone()
              sbbx1, sbby1, sbbx2, sbby2 = (
                  self.pooling_scale * bbx1,
                  self.pooling_scale * bby1,
                  self.pooling_scale * bbx2,
                  self.pooling_scale * bby2,
              )
              temp_x[:, sbbx1:sbbx2, sbby1:sbby2, :] = x.flip(0)[
                  :, sbbx1:sbbx2, sbby1:sbby2, :
              ]
              x = temp_x
          else:
              bbx1, bby1, bbx2, bby2 = 0, 0, 0, 0
  
          x = self.forward_tokens(x, h, w)
          x_cls = self.head(x[:, 0])
          x_aux = self.aux_head(
              x[:, 1:]
          )  # generate classes in all feature tokens, see token labeling
  
          if not self.training:
              return x_cls + 0.5 * x_aux.max(1)[0]
  
          if (
              self.mix_token and self.training
          ):  # reverse "mix token", see token labeling for details.
              x_aux = x_aux.reshape(
                  x_aux.shape[0], patch_h, patch_w, x_aux.shape[-1]
              )
  
              temp_x = x_aux.clone()
              temp_x[:, bbx1:bbx2, bby1:bby2, :] = x_aux.flip(0)[
                  :, bbx1:bbx2, bby1:bby2, :
              ]
              x_aux = temp_x
  
              x_aux = x_aux.reshape(
                  x_aux.shape[0], patch_h * patch_w, x_aux.shape[-1]
              )
  
          return x_cls, x_aux, (bbx1, bby1, bbx2, bby2)
  
file_path: towhee/towhee/models/wave_vit/wave_vit.py
imports: import torch
  from torch import nn
  from typing import List, Any
  import numpy
  from timm.models.vision_transformer import _cfg
  from towhee.models.wave_vit.wave_vit_block import rand_bbox, Stem, DownSamples, Block, ClassBlock
  from towhee.models.utils.init_vit_weights import init_vit_weights
  from towhee.models.wave_vit.wave_vit_utils import get_configs
label_desc: Use image node_id 4 for calling the WaveViT obj's underlying member method code with example usage: obj.forward(x) and returns: x, x_cls,  x_aux,  bbx1,  bby1,  bbx2,  bby2, unknown