projects
/
platform
/
upstream
/
tensorflow.git
/ commitdiff
commit
grep
author
committer
pickaxe
?
search:
re
summary
|
shortlog
|
log
|
commit
| commitdiff |
tree
raw
|
patch
| inline |
side by side
(parent:
7586dee
)
Add shape check for compression_type in TFrecordDataset
author
Yong Tang
<yong.tang.github@outlook.com>
Tue, 17 Apr 2018 01:30:42 +0000
(
01:30
+0000)
committer
Yong Tang
<yong.tang.github@outlook.com>
Tue, 17 Apr 2018 01:30:42 +0000
(
01:30
+0000)
Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
tensorflow/core/ops/dataset_ops.cc
patch
|
blob
|
history
diff --git
a/tensorflow/core/ops/dataset_ops.cc
b/tensorflow/core/ops/dataset_ops.cc
index
47a0c0b
..
ce28a9c
100644
(file)
--- 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);
});