src/transformers/models/esm/modeling_esmfold.py/EsmFoldAngleResnetBlock/forward
class EsmFoldAngleResnetBlock:    def forward(self, a: torch.Tensor) -> torch.Tensor:
        s_initial = a

        a = self.relu(a)
        a = self.linear_1(a)
        a = self.relu(a)
        a = self.linear_2(a)

        return a + s_initial