def METHOD_NAME(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
):
    input_shape = input_ids.size()
    device = input_ids.device
    if attention_mask is None:
        attention_mask = (
            torch.ones(input_shape, device=device)
            if input_ids is None
            else (input_ids != self.bert.config.pad_token_id)
        )
    outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
    sequence_output = outputs.last_hidden_state
    tok_weights = self.tok_proj(sequence_output)
    tok_weights = torch.relu(tok_weights)
    return tok_weights