I am using a “one hot” implementation of Cross Entropy Loss, meaning the target is also a vector and not an index, I need this kind of implementation for further research.
When I compare pytorch nn.CrossEntropyLoss (when giving target as an index instead of “one hot”) to my implementation,I can’t learn anything, I suspect it has to do with vanishing gradients.
Both logits and targets are of type float tensor.
def forward(self, logits, targets):
    return torch.sum(- targets * F.log_softmax(logits, -1), -1)

another implementation I tried and didn’t work well, which should be similar in case of one hot:
def forward(self, logits, targets):
     -1 * torch.log(torch.sum(targets * F.softmax(logits, -1), -1))

weirdly enough, when I change the model architecture to a more complex one, the loss works well.
Any ideas?