Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

fix nll_loss and cross_entropy

This commit is contained in:
lvyufeng 2022-08-27 08:45:42 +08:00
parent 2df40c2deb
commit ff07db1f69
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
mindspore/python/mindspore/ops/function/nn_func.py
View File

@ -1352,6 +1352,11 @@ def _nll_loss(inputs, target, target_dim=-1, weight=None, ignore_index=None, red
if target.ndim == inputs.ndim - 1:
target = target.expand_dims(target_dim)
if ignore_index is not None:
non_pad_mask = _equal(target, ignore_index)
target = target.masked_fill(non_pad_mask, 0)
else:
non_pad_mask = target
loss = _neg(_gather_d(inputs, target_dim, target))
smooth_loss = _neg(inputs.sum(axis=target_dim, keepdims=True))
if weight is not None:
@ -1360,7 +1365,6 @@ def _nll_loss(inputs, target, target_dim=-1, weight=None, ignore_index=None, red
else:
loss_weights = _ones_like(loss)
if ignore_index is not None:
non_pad_mask = _equal(target, ignore_index)
loss = loss.masked_fill(non_pad_mask, 0.)
loss_weights = loss_weights.masked_fill(non_pad_mask, 0.)
smooth_loss = smooth_loss.masked_fill(non_pad_mask, 0.)