def forward(self, s: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            [*, N_res, C_s] single representation
        Returns:
            [*, N_res, 6] update vector
        """
        # [*, 6]
        update = self.linear(s)

        return update

        Args:
            [*, N_res, C_s] single representation
        Returns:
            [*, N_res, 6] update vector
        
dpkd/transformers/src/transformers/models/esm/modeling_esmfold.py