From 53dc0143377552418f1c4db39c5a388a75fd52f8 Mon Sep 17 00:00:00 2001 From: HydrogenSulfate <490868991@qq.com> Date: Mon, 11 Oct 2021 20:36:26 +0800 Subject: [PATCH] Update loss.py --- python/paddle/nn/functional/loss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/paddle/nn/functional/loss.py b/python/paddle/nn/functional/loss.py index f8e03e476d7..5bb317cf3e7 100755 --- a/python/paddle/nn/functional/loss.py +++ b/python/paddle/nn/functional/loss.py @@ -1712,9 +1712,9 @@ def cross_entropy(input, if ignore_weight_mask.ndim > 1 and ignore_weight_mask.shape[ axis] == 1: ignore_weight_mask.squeeze_(axis) - if axis != -1: + if axis != -1 and axis != valid_label.ndim - 1: temp_perm = list(range(axis % valid_label.ndim)) \ - + list(range((axis + 1) % valid_label.ndim, valid_label.ndim)) \ + + list(range((axis % valid_label.ndim + 1) , valid_label.ndim)) \ + [axis%valid_label.ndim] weight_gather = _C_ops.gather_nd( weight, valid_label.transpose(temp_perm)) -- GitLab