def _space_to_depth(self, input):
        """3D space to depth trick for TPU optimization.
      """
        B, C, T, H, W = input.shape
        input = input.view(B, C, T // 2, 2, H // 2, 2, W // 2, 2)
        input = input.permute(0, 3, 5, 7, 1, 2, 4, 6)
        input = input.contiguous().view(B, 8 * C, T // 2, H // 2, W // 2)
        return input
3D space to depth trick for TPU optimization.
      
structured_prompting/fairseq-version/fairseq/examples/MMPT/mmpt/processors/models/s3dg.py