From d97ffbdf362fa7d06ef8d946c8620ff7a3a50a08 Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Tue, 17 Apr 2018 01:30:42 +0000 Subject: [PATCH] Add shape check for compression_type in TFrecordDataset Signed-off-by: Yong Tang --- tensorflow/core/ops/dataset_ops.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorflow/core/ops/dataset_ops.cc b/tensorflow/core/ops/dataset_ops.cc index 47a0c0b..ce28a9c 100644 --- a/tensorflow/core/ops/dataset_ops.cc +++ b/tensorflow/core/ops/dataset_ops.cc @@ -421,6 +421,8 @@ REGISTER_OP("TFRecordDataset") shape_inference::ShapeHandle unused; // `filenames` must be a scalar or a vector. TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused)); + // `compression_type` could only be a scalar. + TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused) ); return shape_inference::ScalarShape(c); }); -- 2.7.4