Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix assertion failure in pooling_ops_3d
PiperOrigin-RevId: 372364504
Change-Id: Iecde4fe26b47a8fa935d6e2611b5585ed5777781
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed May 6, 2021
1 parent d80ffba commit 6fc9141
Showing 1 changed file with 13 additions and 0 deletions.
13 changes: 13 additions & 0 deletions tensorflow/core/kernels/pooling_ops_3d.cc
Expand Up @@ -383,6 +383,19 @@ struct LaunchAvgPooling3dGradOp<CPUDevice, T> {
const std::array<int64, 3>& output_shape,
const std::array<int64, 3>& padding,
TensorFormat data_format, Tensor* output) {
OP_REQUIRES(
context, tensor_in_shape.dim_size(0) == out_backprop.dim_size(0),
errors::InvalidArgument(
"Expected first dimension of tensor_in_shape and "
"out_backprop to match, got ",
tensor_in_shape.dim_size(0), " and ", out_backprop.dim_size(0)));
OP_REQUIRES(
context, tensor_in_shape.dim_size(4) == out_backprop.dim_size(4),
errors::InvalidArgument(
"Expected last dimension of tensor_in_shape and "
"out_backprop to match, got ",
tensor_in_shape.dim_size(4), " and ", out_backprop.dim_size(4)));

output->flat<T>().setZero();
std::array<int64, 3> input_size = {{tensor_in_shape.dim_size(3),
tensor_in_shape.dim_size(2),
Expand Down

0 comments on commit 6fc9141

Please sign in to comment.