diff options
Diffstat (limited to 'tensorflow/core/ops/random_ops.cc')
| -rw-r--r-- | tensorflow/core/ops/random_ops.cc | 7 |
1 files changed, 6 insertions, 1 deletions
diff --git a/tensorflow/core/ops/random_ops.cc b/tensorflow/core/ops/random_ops.cc index f6c668f5c9..416ce9c0d8 100644 --- a/tensorflow/core/ops/random_ops.cc +++ b/tensorflow/core/ops/random_ops.cc @@ -43,7 +43,12 @@ REGISTER_OP("RandomUniformInt") .Attr("seed2: int = 0") .Attr("Tout: {int32, int64}") .Attr("T: {int32, int64}") - .SetShapeFn(shape_inference::RandomShape); + .SetShapeFn([](InferenceContext* c) { + ShapeHandle unused; + TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); + TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); + return shape_inference::RandomShape(c); + }); REGISTER_OP("RandomStandardNormal") .Input("shape: T") |