diff options
Diffstat (limited to 'tensorflow/core/ops/lookup_ops.cc')
| -rw-r--r-- | tensorflow/core/ops/lookup_ops.cc | 2 |
1 files changed, 2 insertions, 0 deletions
diff --git a/tensorflow/core/ops/lookup_ops.cc b/tensorflow/core/ops/lookup_ops.cc index 444aa8b954..2059741da9 100644 --- a/tensorflow/core/ops/lookup_ops.cc +++ b/tensorflow/core/ops/lookup_ops.cc @@ -140,11 +140,13 @@ REGISTER_OP("LookupTableSize") .Input("table_handle: Ref(string)") .Output("size: int64") .SetShapeFn(TwoElementVectorInputsAndScalarOutputs); +WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS("LookupTableSize"); REGISTER_OP("LookupTableSizeV2") .Input("table_handle: resource") .Output("size: int64") .SetShapeFn(ScalarAndTwoElementVectorInputsAndScalarOutputs); +WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS("LookupTableSizeV2"); REGISTER_OP("LookupTableExport") .Input("table_handle: Ref(string)") |