commit c9bc6ad023b42195c3ab0ddb824b974146351ad1
author Jiri Simsa <jsimsa@google.com> Fri Aug 05 12:34:32 2022 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Fri Aug 05 12:36:50 2022 -0700
tree 7dde97490976f4bedbefa4509587831a9b5bcef4
parent 884acc353437e2fca9f0e2f7aabc8d9adc3ebd17

[tf.data] Fix return value type in `tf.data.Dataset.save` documentation.

PiperOrigin-RevId: 465624516

tensorflow/python/data/experimental/ops/io.py

diff --git a/tensorflow/python/data/experimental/ops/io.py b/tensorflow/python/data/experimental/ops/io.py
index 075f16f..9eb3630 100644
--- a/tensorflow/python/data/experimental/ops/io.py
+++ b/tensorflow/python/data/experimental/ops/io.py
@@ -59,7 +59,7 @@
   ```python
   dataset = make_dataset()
   def custom_shard_func(element):
-    return 0
+    return np.int64(0)
   dataset = tf.data.experimental.save(
       path="/path/to/data", ..., shard_func=custom_shard_func)
   ```

tensorflow/python/data/ops/dataset_ops.py

diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py
index 1a32e6a..12a80b1 100644
--- a/tensorflow/python/data/ops/dataset_ops.py
+++ b/tensorflow/python/data/ops/dataset_ops.py
@@ -1721,7 +1721,7 @@
       ```python
       dataset = make_dataset()
       def custom_shard_func(element):
-        return 0
+        return np.int64(0)
       dataset.save(
           path="/path/to/data", ..., shard_func=custom_shard_func)
       ```