tejaskhot

# pytorch function to replicate tensorflow's tf.nn.softmax_cross_entropy_with_logits
# works for soft targets or one-hot encodings

import torch
import torch.nn.functional as F

logits = model(input)
loss = torch.sum(- target * F.log_softmax(logits, -1), -1)
mean_loss = loss.mean()