esm/axial_attention.py [40:46]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    def _batched_forward(
        self,
        x,
        self_attn_mask=None,
        self_attn_padding_mask=None,
    ):
        num_rows, num_cols, batch_size, embed_dim = x.size()
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



esm/axial_attention.py [158:164]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    def _batched_forward(
        self,
        x,
        self_attn_mask=None,
        self_attn_padding_mask=None,
    ):
        num_rows, num_cols, batch_size, embed_dim = x.size()
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -