@justusschock Good point with zero loss not contributing to backprop. In newer versions of PyTorch, nn.CrossEntropyLoss has an ignore_index parameter which applies your logic internally to save a bit of code. For instance, no matter what the pred tensor is, the loss will always be zero if target contains only zeroes
import torch
import torch.nn as nn

criterion = nn.CrossEntropyLoss(ignore_index=0)
pred = torch.randn(3, 5)
target = torch.zeros(3, dtype=torch.long)
criterion(pred, target)