code: def forward(self, x, B, T, W):
      num_spatial_tokens = (x.size(1) - 1) // T
      H = num_spatial_tokens // W
  
      if self.attention_type in ["space_only", "joint_space_time"]:
          x = x + self.drop_path(self.attn(self.norm1(x)))
          x = x + self.drop_path(self.mlp(self.norm2(x)))
          return x
      elif self.attention_type == "divided_space_time":
          # Temporal
          xt = x[:, 1:, :]
          xt = rearrange(
              xt, "b (h w t) m -> (b h w) t m", b=B, h=H, w=W, t=T
          )
  
          temporal_attn_out = self.temporal_attn(
              self.temporal_norm1(xt)
          )
  
          res_temporal = self.drop_path(temporal_attn_out)
  
          res_temporal = rearrange(
              res_temporal,
              "(b h w) t m -> b (h w t) m",
              b=B,
              h=H,
              w=W,
              t=T,
          )
          res_temporal = self.temporal_fc(res_temporal)
          xt = x[:, 1:, :] + res_temporal
  
          # Spatial
          init_cls_token = x[:, 0, :].unsqueeze(1)
          cls_token = init_cls_token.repeat(1, T, 1)
          cls_token = rearrange(
              cls_token, "b t m -> (b t) m", b=B, t=T
          ).unsqueeze(1)
          xs = xt
          xs = rearrange(
              xs, "b (h w t) m -> (b t) (h w) m", b=B, h=H, w=W, t=T
          )
          xs = torch.cat((cls_token, xs), 1)
  
          spatial_attn_out = self.attn(self.norm1(xs))
          res_spatial = self.drop_path(spatial_attn_out)
  
          # Taking care of CLS token
          cls_token = res_spatial[:, 0, :]
          cls_token = rearrange(cls_token, "(b t) m -> b t m", b=B, t=T)
          # averaging for every frame
          cls_token = torch.mean(cls_token, 1, True)
          res_spatial = res_spatial[:, 1:, :]
          res_spatial = rearrange(
              res_spatial,
              "(b t) (h w) m -> b (h w t) m",
              b=B,
              h=H,
              w=W,
              t=T,
          )
          res = res_spatial
          x = xt
  
          # Mlp
          x = torch.cat((init_cls_token, x), 1) + torch.cat(
              (cls_token, res), 1
          )
  
          x_res = x
  
          x = self.norm2(x)
          # x = x + self.drop_path(self.mlp(self.norm2(x)))
  
          # MLP
          mlp_out = self.mlp(x)
  
          x = x_res + self.drop_path(mlp_out)
          return x
  
imports: import logging
  from functools import partial
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils
  import torch.utils.checkpoint
  from einops import rearrange
  from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
  from .helpers import load_pretrained, load_pretrained_imagenet, load_pretrained_kinetics
  from .vit_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, DropPath, to_2tuple, trunc_normal_
label_desc: Use image node_id 2 for calling the Block obj's underlying member method code with example usage: obj.forward(x, B, T, W) and returns: x, x