def forward(
            self,
            hidden_states: torch.Tensor,
            attention_mask: torch.Tensor,
            causal_attention_mask: torch.Tensor,
            output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.FloatTensor]:
        """
        Args:
            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
            attention_mask (`torch.FloatTensor`): attention mask of size
                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
                `(config.encoder_attention_heads,)`.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
        """

        if self.add_time_attn:
            bt, n, d = hidden_states.shape
            t = self.t

            # time embed
            if t != 1:
                n = hidden_states.shape[1]
                hidden_states = rearrange(hidden_states, '(b t) n d -> (b n) t d', t=t)
                hidden_states = hidden_states + self.temporal_embedding[:, :t, :]
                hidden_states = rearrange(hidden_states, '(b n) t d -> (b t) n d', n=n)

            # time attn
            residual = hidden_states
            hidden_states = rearrange(hidden_states, '(b t) n d -> (b n) t d', t=t)
            # hidden_states = self.layer_norm1(hidden_states)  # share layernorm
            hidden_states = self.temporal_layer_norm1(hidden_states)
            hidden_states, attn_weights = self.temporal_attn(
                hidden_states=hidden_states,
                attention_mask=attention_mask,
                causal_attention_mask=causal_attention_mask,
                output_attentions=output_attentions,
            )
            hidden_states = residual + rearrange(hidden_states, '(b n) t d -> (b t) n d', n=n)

            # residual = hidden_states
            # hidden_states = rearrange(hidden_states, '(b t) n d -> (b n) t d', t=t)
            # # hidden_states = self.layer_norm2(hidden_states)  # share layernorm
            # hidden_states = self.temporal_layer_norm2(hidden_states)
            # hidden_states = self.temporal_mlp(hidden_states)
            # hidden_states = residual + rearrange(hidden_states, '(b n) t d -> (b t) n d', n=n)

        # spatial attn
        residual = hidden_states

        hidden_states = self.layer_norm1(hidden_states)
        hidden_states, attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            causal_attention_mask=causal_attention_mask,
            output_attentions=output_attentions,
        )
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.layer_norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (attn_weights,)

        return outputs

        Args:
            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
            attention_mask (`torch.FloatTensor`): attention mask of size
                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
                `(config.encoder_attention_heads,)`.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
        
src/marqo/s2_inference/languagebind/video/modeling_video.py