def forward(self, hidden_states: torch.Tensor):
        """
        Args:
            hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
        """
        hidden_states = self.w_in(hidden_states)

        if self.dropout is not None:
            hidden_states = self.dropout(hidden_states)

        hidden_states = self.w_out(hidden_states)

        return hidden_states

        Args:
            hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
        
dpkd/transformers/src/transformers/models/cpmant/modeling_cpmant.py