Check label range in cross entropy calculation. (#10954)

7 years ago · 3ba75d4a69
parent 0c44efb855
commit 3ba75d4a69
1 changed files with 4 additions and 1 deletions
--- a/paddle/fluid/operators/math/cross_entropy.cc
+++ b/paddle/fluid/operators/math/cross_entropy.cc
@ -46,7 +46,10 @@ class CrossEntropyFunctor<platform::CPUDeviceContext, T> {

      const int64_t* label_data = labels->data<int64_t>();
      for (int i = 0; i < batch_size; ++i) {
-        int index = i * class_num + label_data[i];
+        int lbl = label_data[i];
+        PADDLE_ENFORCE_GE(lbl, 0);
+        PADDLE_ENFORCE_LT(lbl, class_num);
+        int index = i * class_num + lbl;
        loss_data[i] = -math::TolerableValue<T>()(std::log(prob_data[index]));
      }
    }