36 lines
1.7 KiB
Diff
36 lines
1.7 KiB
Diff
From 6fc9141f42f6a72180ecd24021c3e6b36165fe0d Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
Date: Thu, 6 May 2021 09:51:26 -0700
|
|
Subject: [PATCH] Fix assertion failure in pooling_ops_3d
|
|
|
|
PiperOrigin-RevId: 372364504
|
|
Change-Id: Iecde4fe26b47a8fa935d6e2611b5585ed5777781
|
|
---
|
|
tensorflow/core/kernels/pooling_ops_3d.cc | 13 +++++++++++++
|
|
1 file changed, 13 insertions(+)
|
|
|
|
diff --git a/tensorflow/core/kernels/pooling_ops_3d.cc b/tensorflow/core/kernels/pooling_ops_3d.cc
|
|
index 9da2d62b0a21d..56a55bc2ec87b 100644
|
|
--- a/tensorflow/core/kernels/pooling_ops_3d.cc
|
|
+++ b/tensorflow/core/kernels/pooling_ops_3d.cc
|
|
@@ -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),
|