    def forward(self, x, mask_matrix):
        """ Forward function.
        Args:
            x (`tensor`):
                Input feature, tensor size (B, D, H, W, C).
            mask_matrix (`tuple[int]`):
                Attention mask for cyclic shift.
        """

        shortcut = x
        if self.use_checkpoint:
            x = checkpoint.checkpoint(self.forward_part1, x, mask_matrix)
        else:
            x = self.forward_part1(x, mask_matrix)
        x = shortcut + self.drop_path(x)

        if self.use_checkpoint:
            x = x + checkpoint.checkpoint(self.forward_part2, x)
        else:
            x = x + self.forward_part2(x)

        return x
Forward function.
Args:
x (`tensor`):
Input feature, tensor size (B, D, H, W, C).
mask_matrix (`tuple[int]`):
Attention mask for cyclic shift.
towhee/models/layers/swin_transformer_block3d.py