Pytorch softmax cross entropy with logits

softmax_cross_entropy_with_logits.py

# 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()

exm, it seems you want to replicate tensorflow's tf.nn.soft,ax_cross_entropy_with_logits, but you use F.log_softmax rather than F.softmax?

exm, it seems you want to replicate tensorflow's tf.nn.soft,ax_cross_entropy_with_logits, but you use F.log_softmax rather than F.softmax?

Here log is for computing cross entropy.

Will this work if my logits shape is (32, 1, 128, 128) and target is also of same shape but with all values zero and one [x,y] = 1.0