Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix FPE issue with tf.raw_ops.FusedBatchNorm.
PiperOrigin-RevId: 370948185
Change-Id: If0c8e0320062ed6363e94ff5fe38e6a301f69ac2
  • Loading branch information
Amit Patankar authored and tensorflower-gardener committed Apr 28, 2021
1 parent 828f346 commit 1a2a872
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions tensorflow/core/kernels/fused_batch_norm_op.cc
Expand Up @@ -293,6 +293,9 @@ struct FusedBatchNorm<CPUDevice, T, U, /* is_training= */ false> {
const CPUDevice& d = context->eigen_device<CPUDevice>();

const int depth = x.dimension(3);
OP_REQUIRES(
context, depth != 0,
errors::Internal("The 4th element in the input shape cannot be 0."));
const int size = x.size();
const int rest_size = size / depth;
Eigen::DSizes<Eigen::Index, 2> rest_by_depth(rest_size, depth);
Expand Down

0 comments on commit 1a2a872

Please sign in to comment.