<p>@Blade  Here's the solution I came up with!</p>

<pre><code>import torch
import torch.nn as nn
import torch.nn.functional as F


class masked_softmax_cross_entropy_loss(nn.Module):
    r"""my version of masked tf.nn.softmax_cross_entropy_with_logits"""
    def __init__(self, weight=None):
        super(masked_softmax_cross_entropy_loss, self).__init__()
        self.register_buffer('weight', weight)

    def forward(self, input, target, mask):
        if not target.is_same_size(input):
            raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))

        input = F.softmax(input)
        loss = -torch.sum(target * torch.log(input), 1)
        loss = torch.unsqueeze(loss, 1)
        mask /= torch.mean(mask)
        mask = torch.unsqueeze(mask, 1)
        loss = torch.mul(loss, mask)
        return torch.mean(loss)
</code></pre>

<p>Btw: I needed this loss function at the time (Sept 2017) because I was attempting to translate Thomas Kipf's GCN (see <a href="https://arxiv.org/abs/1609.02907" rel="nofollow noreferrer">https://arxiv.org/abs/1609.02907</a>) code from TensorFlow to PyTorch.  However, I now notice that Kipf has done this himself (see <a href="https://github.com/tkipf/pygcn" rel="nofollow noreferrer">https://github.com/tkipf/pygcn</a>), and in his code, he simply uses the built-in PyTorch loss function, the negative log likelihood loss, i.e.</p>

<pre><code>loss_train = F.nll_loss(output[idx_train], labels[idx_train])
</code></pre>

<p>Hope this helps.</p>

<p>~DV</p>
PyTorch equivalence for softmax_cross_entropy_with_logits
<p>I was wondering is there an equivalent PyTorch loss function for TensorFlow's <code>softmax_cross_entropy_with_logits</code>?</p>