src/transformers/models/esm/modeling_esmfold.py/EsmFoldPairToSequence/__init__
class EsmFoldPairToSequence:    def __init__(self, pairwise_state_dim, num_heads):
        super().__init__()

        self.layernorm = nn.LayerNorm(pairwise_state_dim)
        self.linear = nn.Linear(pairwise_state_dim, num_heads, bias=False)