    def forward(self, embeddings: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """
        # Parameters

        embeddings : `torch.Tensor`
            A tensor of shape `(sequence_length, embedding_dim)`
        targets : `torch.Tensor`
            A tensor of shape `(batch_size, )`

        # Returns

        loss : `torch.FloatTensor`
            A scalar loss to be optimized.
        """
        # embeddings is size (n, embedding_dim)
        # targets is (batch_size, ) with the correct class id
        # Does not do any count normalization / divide by batch size
        probs = torch.nn.functional.log_softmax(
            torch.matmul(embeddings, self.softmax_w) + self.softmax_b, dim=-1
        )

        return torch.nn.functional.nll_loss(probs, targets.long(), reduction="sum")
# Parameters

embeddings : `torch.Tensor`
A tensor of shape `(sequence_length, embedding_dim)`
targets : `torch.Tensor`
A tensor of shape `(batch_size, )`

# Returns

loss : `torch.FloatTensor`
A scalar loss to be optimized.
allennlp/modules/softmax_loss.py