diff --git a/orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc b/orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc new file mode 100644 index 0000000000000..8f72bdccdd589 --- /dev/null +++ b/orttraining/orttraining/test/training_ops/cpu/loss/cross_entropy_test.cc @@ -0,0 +1,214 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "gtest/gtest.h" +#include "test/providers/provider_test_utils.h" + +namespace onnxruntime { +namespace test { + +// Regression tests for OOB reads when label values are outside [0, C). + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLarge) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, 5, 2}; // 5 is out of range [0, 5) + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_NegativeLabel) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-100)); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, -1, 2}; // -1 is out of range (and != ignore_index) + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLargeWithWeights) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, 100, 2}; // 100 is out of range + std::vector weight_data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f}; + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + test.AddInput("weight", {5}, weight_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +// Covers the weighted, non-MEAN forward loop (the second per-sample loop in Compute). +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLargeWithWeightsSumReduction) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("sum")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, 7, 2}; // 7 is out of range [0, 5) + std::vector weight_data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f}; + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + test.AddInput("weight", {5}, weight_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +// int32 label type — kernel is registered for both int32_t and int64_t. +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLargeInt32) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, 5, 2}; // 5 is out of range [0, 5) + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +// Higher-dimensional inputs: logit [N, C, D1, D2], label [N, D1, D2]. +TEST(CrossEntropyTest, SoftmaxCrossEntropyLoss_LabelTooLargeHighDim) { + OpTester test("SoftmaxCrossEntropyLoss", 12); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + // [N=2, C=4, D1=2, D2=3] -> label shape [2, 2, 3] -> 12 label entries. + std::vector X_data(2 * 4 * 2 * 3, 1.0f); + std::vector index_data = {0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 99, 3}; // 99 is out of range + std::vector log_prob_init(2 * 4 * 2 * 3, 0.0f); + + test.AddInput("X", {2, 4, 2, 3}, X_data); + test.AddInput("index", {2, 2, 3}, index_data); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {2, 4, 2, 3}, log_prob_init); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LabelTooLarge) { + OpTester test("SoftmaxCrossEntropyLossGrad", 1, onnxruntime::kMSDomain); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector dY_data = {1.0f}; + std::vector log_prob_data(3 * 5, -1.6094f); + std::vector index_data = {0, 5, 2}; // 5 is out of range [0, 5) + + test.AddInput("dY", {}, dY_data); + test.AddInput("log_prob", {3, 5}, log_prob_data); + test.AddInput("index", {3}, index_data); + test.AddOutput("dX", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LabelTooLargeWithWeights) { + OpTester test("SoftmaxCrossEntropyLossGrad", 1, onnxruntime::kMSDomain); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector dY_data = {1.0f}; + std::vector log_prob_data(3 * 5, -1.6094f); + std::vector index_data = {0, 5, 2}; // 5 is out of range [0, 5) + std::vector weight_data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f}; + + test.AddInput("dY", {}, dY_data); + test.AddInput("log_prob", {3, 5}, log_prob_data); + test.AddInput("index", {3}, index_data); + test.AddInput("weight", {5}, weight_data); + test.AddOutput("dX", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossGrad_LabelTooLargeInt32) { + OpTester test("SoftmaxCrossEntropyLossGrad", 1, onnxruntime::kMSDomain); + test.AddAttribute("reduction", std::string("mean")); + test.AddAttribute("ignore_index", static_cast(-1)); + + std::vector dY_data = {1.0f}; + std::vector log_prob_data(3 * 5, -1.6094f); + std::vector index_data = {0, 5, 2}; + + test.AddInput("dY", {}, dY_data); + test.AddInput("log_prob", {3, 5}, log_prob_data); + test.AddInput("index", {3}, index_data); + test.AddOutput("dX", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +// SoftmaxCrossEntropyLossInternal shares the same Compute as SoftmaxCrossEntropyLoss but +// is registered separately under kMSDomain v1 with an optional runtime ignore_index input. +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossInternal_LabelTooLarge) { + OpTester test("SoftmaxCrossEntropyLossInternal", 1, onnxruntime::kMSDomain); + test.AddAttribute("reduction", std::string("mean")); + + std::vector X_data(3 * 5, 1.0f); + std::vector index_data = {0, 5, 2}; + int64_t ignore_index_val = -1; + + test.AddInput("X", {3, 5}, X_data); + test.AddInput("index", {3}, index_data); + // weight is optional; pass empty input to skip and still provide ignore_index input below. + test.AddOptionalInputEdge(); + test.AddInput("ignore_index", {}, &ignore_index_val, 1); + test.AddOutput("output", {}, {0.0f}); + test.AddOutput("log_prob", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +TEST(CrossEntropyTest, SoftmaxCrossEntropyLossInternalGrad_LabelTooLarge) { + OpTester test("SoftmaxCrossEntropyLossInternalGrad", 1, onnxruntime::kMSDomain); + test.AddAttribute("reduction", std::string("mean")); + + std::vector dY_data = {1.0f}; + std::vector log_prob_data(3 * 5, -1.6094f); + std::vector index_data = {0, 5, 2}; + int64_t ignore_index_val = -1; + + test.AddInput("dY", {}, dY_data); + test.AddInput("log_prob", {3, 5}, log_prob_data); + test.AddInput("index", {3}, index_data); + test.AddOptionalInputEdge(); // weight + test.AddInput("ignore_index", {}, &ignore_index_val, 1); + test.AddOutput("dX", {3, 5}, std::vector(15, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "out of range"); +} + +} // namespace test +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc index c74bf06a77d6e..b575344219d47 100644 --- a/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc +++ b/orttraining/orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.cc @@ -10,6 +10,8 @@ #include "core/providers/cpu/controlflow/scan_utils.h" #include "orttraining/training_ops/cpu/loss/cross_entropy.h" #include "orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.h" +#include + #include namespace onnxruntime { @@ -56,12 +58,17 @@ void GetNDCFromLogitAndLabelShape(const TensorShape& logit_shape, const TensorSh void VerifyLogitWeightAndLabelShape(const TensorShape& logit_shape, const TensorShape& label_shape, const TensorShape* weight_shape) { - ORT_ENFORCE(nullptr == weight_shape || 1 == weight_shape->NumDimensions(), "Weights tensor is not 1-D."); - const size_t label_dims = label_shape.NumDimensions(); + ORT_ENFORCE(label_dims >= 1, "label must be at least 1-D."); + ORT_ENFORCE(logit_shape.NumDimensions() >= 2, "logit must be at least 2-D."); ORT_ENFORCE(logit_shape.NumDimensions() == label_dims + 1, "logit_shape must be (1 + label_shape)"); + ORT_ENFORCE(nullptr == weight_shape || 1 == weight_shape->NumDimensions(), "Weights tensor is not 1-D."); + ORT_ENFORCE(nullptr == weight_shape || (*weight_shape)[0] == logit_shape[1], + "Weight tensor size (", (weight_shape ? (*weight_shape)[0] : 0), + ") must equal the number of classes (", logit_shape[1], ")"); + ORT_ENFORCE(label_shape[0] == logit_shape[0], "The shape of logit and label does not match"); if (label_dims >= 2) { @@ -103,7 +110,7 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const const Tensor* p_ignore_index = context->Input(3); int64_t ignore_index = ignore_index_; if (p_ignore_index) { - ORT_ENFORCE(p_ignore_index->Shape().IsScalar(), "ignore_index should be a scalar."); + ORT_RETURN_IF_NOT(p_ignore_index->Shape().IsScalar(), "ignore_index should be a scalar."); ignore_index = *(p_ignore_index->template Data()); } @@ -111,6 +118,12 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const const TensorShape label_shape{label.Shape()}; VerifyLogitWeightAndLabelShape(logit_shape, label_shape, p_weight ? &p_weight->Shape() : nullptr); + // Pre-empt the divide-by-zero in GetNDCFromLogitAndLabelShape (which does + // C = logit_shape.Size() / N_D). VerifyLogitWeightAndLabelShape only checks + // rank, not Size > 0 — e.g. label shape [3, 0, 5] is rank 2 but empty. + ORT_RETURN_IF_NOT(label_shape.Size() > 0, + "SoftmaxCrossEntropyLoss: label tensor must not be empty."); + // N_D = N * D1 * D2...D*K int64_t N_D = 0; int64_t C = 0; @@ -130,9 +143,27 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const logit_data = (*transpose_output.GetMutable()).template Data(); } + // Validate N_D, C and compute N_D * C with overflow checks. + // Uses the no-exceptions-safe pattern from deform_conv_attributes.h: explicit + // pre-multiply guards via ORT_RETURN_IF_NOT, then a plain multiply, then narrow_cast. + constexpr int64_t kInt64Max = std::numeric_limits::max(); + constexpr int64_t kIntMax = static_cast(std::numeric_limits::max()); + constexpr int64_t kEigenIndexMax = static_cast(std::numeric_limits::max()); + + ORT_RETURN_IF_NOT(N_D > 0 && C > 0, + "SoftmaxCrossEntropyLoss: N_D and C must be positive (got N_D=", N_D, ", C=", C, ")."); + ORT_RETURN_IF_NOT(N_D <= kIntMax && C <= kIntMax, + "SoftmaxCrossEntropyLoss: N_D=", N_D, ", C=", C, " exceed int max."); + ORT_RETURN_IF_NOT(N_D <= kInt64Max / C, + "SoftmaxCrossEntropyLoss: N_D * C overflows int64 (N_D=", N_D, ", C=", C, ")."); + const int64_t n_d_c_i64 = N_D * C; + ORT_RETURN_IF_NOT(n_d_c_i64 <= kEigenIndexMax, + "SoftmaxCrossEntropyLoss: N_D * C (", n_d_c_i64, ") exceeds Eigen::Index max."); + const int n_d = gsl::narrow_cast(N_D); const int c = gsl::narrow_cast(C); - const uint64_t n_d_c = N_D * C; + const uint64_t n_d_c = static_cast(n_d_c_i64); + Tensor* loss = context->Output(0, reduction_ == ReductionType::NONE ? TensorShape(label.Shape()) : TensorShape({})); T1* log_prob_data; std::vector log_prob_data_buffer(0); @@ -147,9 +178,9 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const } const T2* label_data = label.template Data(); + T1* loss_data = loss->template MutableData(); std::vector shifted_logit(narrow(n_d_c)); - ORT_ENFORCE(n_d_c <= static_cast(std::numeric_limits::max())); ComputeShareSoftmaxCrossEntropyCPU(n_d, c, static_cast(n_d_c), logit_data, shifted_logit.data(), log_prob_data); std::vector loss_sample_buffer(0); @@ -175,19 +206,31 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const // Compute weighed loss for each sample while summing weights for unignored target/label values. if (reduction_ == ReductionType::MEAN) { for (ptrdiff_t i = 0; i < n_d; i++) { - if (ignore_index == label_data[i]) { + const T2 label_sample = label_data[i]; + if (ignore_index == label_sample) { loss_sample[i] = 0; } else { - loss_sample[i] = -log_prob_data[i * c + label_data[i]] * weight_data[label_data[i]]; - sum_weight += weight_data[label_data[i]]; + if (label_sample < 0 || label_sample >= C) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "SoftmaxCrossEntropyLoss: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); + } + loss_sample[i] = -log_prob_data[i * c + label_sample] * weight_data[label_sample]; + sum_weight += weight_data[label_sample]; } } } else { for (ptrdiff_t i = 0; i < n_d; i++) { - if (ignore_index == label_data[i]) { + const T2 label_sample = label_data[i]; + if (ignore_index == label_sample) { loss_sample[i] = 0; } else { - loss_sample[i] = -log_prob_data[i * c + label_data[i]] * weight_data[label_data[i]]; + if (label_sample < 0 || label_sample >= C) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "SoftmaxCrossEntropyLoss: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); + } + loss_sample[i] = -log_prob_data[i * c + label_sample] * weight_data[label_sample]; } } } @@ -205,10 +248,16 @@ Status SoftmaxCrossEntropyLoss::Compute(OpKernelContext* context) const // Compute loss for each sample while counting unignored target/label values. int unignored_samples = 0; for (ptrdiff_t i = 0; i < n_d; i++) { - if (ignore_index == label_data[i]) { + const T2 label_sample = label_data[i]; + if (ignore_index == label_sample) { loss_sample[i] = 0; } else { - loss_sample[i] = -log_prob_data[i * c + label_data[i]]; + if (label_sample < 0 || label_sample >= C) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "SoftmaxCrossEntropyLoss: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); + } + loss_sample[i] = -log_prob_data[i * c + label_sample]; unignored_samples += 1; } } @@ -252,7 +301,7 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co const Tensor* p_ignore_index = context->Input(4); int64_t ignore_index = ignore_index_; if (p_ignore_index) { - ORT_ENFORCE(p_ignore_index->Shape().IsScalar(), "ignore_index should be a scalar."); + ORT_RETURN_IF_NOT(p_ignore_index->Shape().IsScalar(), "ignore_index should be a scalar."); ignore_index = *(p_ignore_index->template Data()); } @@ -260,16 +309,51 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co const TensorShape label_shape{label.Shape()}; VerifyLogitWeightAndLabelShape(probability_shape, label_shape, p_weight ? &p_weight->Shape() : nullptr); + // Pre-empt the divide-by-zero in GetNDCFromLogitAndLabelShape (see forward Compute). + ORT_RETURN_IF_NOT(label_shape.Size() > 0, + "SoftmaxCrossEntropyLossGrad: label tensor must not be empty."); + // N_D = N * D1 * D2...D*K int64_t N_D = 0; int64_t C = 0; GetNDCFromLogitAndLabelShape(probability_shape, label_shape, N_D, C); + + // Compute N_D * C once with overflow checks; reused by every parallel-for below. + // Backward keeps N_D and C as int64_t in lambdas, so only the product needs to fit. + // TryParallelFor takes std::ptrdiff_t, which is 32-bit on 32-bit minimal builds — + // use ptrdiff_t max as the ceiling, not int64_t max. + constexpr int64_t kInt64Max = std::numeric_limits::max(); + constexpr int64_t kPtrdiffMax = static_cast(std::numeric_limits::max()); + + ORT_RETURN_IF_NOT(N_D > 0 && C > 0, + "SoftmaxCrossEntropyLossGrad: N_D and C must be positive (got N_D=", N_D, ", C=", C, ")."); + ORT_RETURN_IF_NOT(N_D <= kInt64Max / C, + "SoftmaxCrossEntropyLossGrad: N_D * C overflows int64 (N_D=", N_D, ", C=", C, ")."); + const int64_t n_d_c_i64 = N_D * C; + ORT_RETURN_IF_NOT(n_d_c_i64 <= kPtrdiffMax, + "SoftmaxCrossEntropyLossGrad: N_D * C (", n_d_c_i64, ") exceeds ptrdiff_t max."); + const std::ptrdiff_t n_d_c = static_cast(n_d_c_i64); + const T1* dY_data = dY.template Data(); const T1* log_prob_data = log_prob.template Data(); const T2* label_data = label.template Data(); + + // Validate label values are within [0, C) to prevent out-of-bounds reads. + // Done as a single up-front O(N_D) pass: returning Status from inside a + // ThreadPool::TryParallelFor lambda would only return from the lambda, not from + // Compute. Cost is negligible vs the O(N_D * C) parallel gradient that follows. + for (int64_t i = 0; i < N_D; i++) { + const T2 label_sample = label_data[i]; + if (ignore_index != label_sample && (label_sample < 0 || label_sample >= C)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "SoftmaxCrossEntropyLossGrad: label value ", label_sample, + " at index ", i, " is out of range [0, ", C, ")"); + } + } + Tensor* d_logit = context->Output(0, probability_shape); T1* d_logit_data = d_logit->template MutableData(); - std::memset(d_logit_data, 0, narrow(sizeof(T1) * N_D)); + std::memset(d_logit_data, 0, narrow(sizeof(T1) * probability_shape.Size())); OrtValue transpose_output; TensorShapeVector new_shape; std::vector permutations; @@ -292,18 +376,18 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co if (reduction_ == ReductionType::NONE) { concurrency::ThreadPool::TryParallelFor( - tp, narrow(N_D * C), cost, + tp, n_d_c, cost, [&label_data, &weight_data, &d_logit_data, &log_prob_data, ignore_index, C, &dY_data]( std::ptrdiff_t begin, std::ptrdiff_t end) { for (std::ptrdiff_t index = begin; index != end; ++index) { int64_t row = index / C; int64_t col = index % C; T2 label_sample = label_data[row]; - T1 weight_smaple = weight_data[label_sample] * dY_data[row]; if (ignore_index == label_sample) { d_logit_data[index] = 0; } else { - d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * weight_smaple; + T1 weight_sample = weight_data[label_sample] * dY_data[row]; + d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * weight_sample; } } }); @@ -323,18 +407,18 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co } concurrency::ThreadPool::TryParallelFor( - tp, narrow(N_D * C), cost, + tp, n_d_c, cost, [&label_data, &weight_data, dY_scaled, &d_logit_data, &log_prob_data, ignore_index, C]( std::ptrdiff_t begin, std::ptrdiff_t end) { for (std::ptrdiff_t index = begin; index != end; ++index) { int64_t row = index / C; int64_t col = index % C; T2 label_sample = label_data[row]; - T1 weight_smaple = weight_data[label_sample] * dY_scaled; if (ignore_index == label_sample) { d_logit_data[index] = 0; } else { - d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * weight_smaple; + T1 weight_sample = weight_data[label_sample] * dY_scaled; + d_logit_data[index] = (exp(log_prob_data[index]) - (label_sample == col)) * weight_sample; } } }); @@ -342,7 +426,7 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co } else { if (reduction_ == ReductionType::NONE) { concurrency::ThreadPool::TryParallelFor( - tp, narrow(N_D * C), cost, + tp, n_d_c, cost, [&label_data, &d_logit_data, &log_prob_data, ignore_index, C, &dY_data]( std::ptrdiff_t begin, std::ptrdiff_t end) { for (std::ptrdiff_t index = begin; index != end; ++index) { @@ -371,7 +455,7 @@ Status SoftmaxCrossEntropyLossGrad::Compute(OpKernelContext* context) co } concurrency::ThreadPool::TryParallelFor( - tp, narrow(N_D * C), cost, + tp, n_d_c, cost, [&label_data, &d_logit_data, &log_prob_data, ignore_index, C, &dY_scaled]( std::ptrdiff_t begin, std::ptrdiff_t end) { for (std::ptrdiff_t index = begin; index != end; ++index) {