def forward(self, x):
        """forward function of Criss-Cross Attention.

        Args:
            x (Tensor): Input feature. \
                shape (batch_size, in_channels, height, width)
        Returns:
            Tensor: Output of the layer, with shape of \
            (batch_size, in_channels, height, width)
        """
        B, C, H, W = x.size()
        query = self.query_conv(x)
        key = self.key_conv(x)
        value = self.value_conv(x)
        energy_H = torch.einsum('bchw,bciw->bwhi', query, key) + NEG_INF_DIAG(
            H, query.device)
        energy_H = energy_H.transpose(1, 2)
        energy_W = torch.einsum('bchw,bchj->bhwj', query, key)
        attn = F.softmax(
            torch.cat([energy_H, energy_W], dim=-1), dim=-1)  # [B,H,W,(H+W)]
        out = torch.einsum('bciw,bhwi->bchw', value, attn[..., :H])
        out += torch.einsum('bchj,bhwj->bchw', value, attn[..., H:])

        out = self.gamma(out) + x
        out = out.contiguous()

        return out
forward function of Criss-Cross Attention.

        Args:
            x (Tensor): Input feature.                 shape (batch_size, in_channels, height, width)
        Returns:
            Tensor: Output of the layer, with shape of             (batch_size, in_channels, height, width)
        
ControlNet-v1-1-nightly/annotator/uniformer/mmcv/ops/cc_attention.py