From a296b167199783770d9a233ceb31f7e4c8d8fe40 Mon Sep 17 00:00:00 2001 From: "M. Zeeshan Siddiqui" Date: Sat, 16 May 2020 00:33:25 -0700 Subject: [PATCH] Prevent divide by zero in CUDA implementation of SoftmaxCrossEntropyLossGrad. (#3962) --- .../cuda/loss/softmax_cross_entropy_loss_impl.cu | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu index 40765d8c94..5f93231c9b 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu @@ -117,7 +117,11 @@ __global__ void _WeightedSoftmaxCrossEntropyLossGrad( int row = i / C; int d = i % C; CUDA_KERNEL_ASSERT(weight[row] == 0 || (label[row] >= 0 && label[row] < C)); - output_data[i] = (*dY) * weight[row] * (_Exp(log_prob[i]) - 1.0 * (d == label[row])) / (*normalize_factor); + if(0 == *normalize_factor){ + output_data[i] = 0; + } else { + output_data[i] = (*dY) * weight[row] * (_Exp(log_prob[i]) - 1.0 * (d == label[row])) / (*normalize_factor); + } } template @@ -135,7 +139,11 @@ __global__ void _WeightedReductionNoneSoftmaxCrossEntropyLossGrad( int row = i / C; int d = i % C; CUDA_KERNEL_ASSERT(weight[row] == 0 || (label[row] >= 0 && label[row] < C)); - output_data[i] = dY[row] * weight[row] * (_Exp(log_prob[i]) - 1.0 * (d == label[row])) / (*normalize_factor); + if(0 == *normalize_factor){ + output_data[i] = 0; + } else { + output_data[i] = dY[row] * weight[row] * (_Exp(log_prob[i]) - 1.0 * (d == label[row])) / (*normalize_factor); + } } template