class_docstr: SplitSABlock
class_docstr_tok: ["SplitSABlock"]
code: def forward(self, x):
      x = x + self.pos_embed(x)
      b, c, t, h, w = x.shape
      attn = x.view(b, c, t, h * w).permute(0, 3, 2, 1).contiguous()
      attn = attn.view(b * h * w, t, c)
      attn = attn + self.drop_path(self.t_attn(self.t_norm(attn)))
      attn = attn.view(b, h * w, t, c).permute(0, 2, 1, 3).contiguous()
      attn = attn.view(b * t, h * w, c)
      residual = x.view(b, c, t, h * w).permute(0, 2, 3, 1).contiguous()
      residual = residual.view(b * t, h * w, c)
      attn = residual + self.drop_path(self.attn(self.norm1(attn)))
      attn = attn.view(b, t * h * w, c)
      out = attn + self.drop_path(self.mlp(self.norm2(attn)))
      out = out.transpose(1, 2).reshape(b, c, t, h, w)
      return out
  
imports: from collections import OrderedDict
  import torch
  from torch import nn
  from torch.utils import checkpoint
  from functools import partial
  from towhee.models.utils.general_utils import to_2tuple
  from towhee.models.utils.weight_init import trunc_normal_
  from towhee.models.layers.droppath import DropPath
  from towhee.models.layers.mlp import Mlp
  from towhee.models.layers.attention import MultiHeadAttention
  from towhee.models.uniformer.config import _C
  import os
label_desc: Use image node_id 2 for calling the SplitSABlock obj's underlying member method code with example usage: obj.forward(x) and returns: out