tf.data.Dataset.save and load custom dataset throw "DataLossError: Unable to parse tensor from stored proto"
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

2.16.1

### Custom code

No

### OS platform and distribution

Linux Ubuntu 22.04

### Mobile device

_No response_

### Python version

3.10.14

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

I am processing a custom dataset from [link](https://colab.research.google.com/drive/1b4PPH4XGht4Jve2xPKMCh-AXXAQziNQa?usp=sharing#scrollTo=HuiDoQnHcNsE). The dataset is in google drive. A sample can be downloaded with `gsutil -m cp -r gs://gresearch/robotics/droid_100 <your_local_path>` and then loaded with `ds = tfds.load("droid_100", data_dir=<your_local_path>, split="train")`

The full data is too large for a single compute to handle. Thus, I need split the dataset into smaller batches, save to local drive and reload each batch for processing. I found no save API in tfds, but I find tf.data.Dataset.save() and load() promising. Unfortunately, after saving and reloading the same dataset, I am getting errors (see relevant log output)

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import tensorflow_datasets as tfds

data_dir = <your_local_path>
ds = tfds.load("droid_100", data_dir=data_dir, split="train").take(1)
ds.save("/tmp/copycat")
ds2 = tf.data.Dataset.load('/tmp/copycat')
for x in ds2:
    pass
```


### Relevant log output

```shell
2024-06-12 07:09:42.851470: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-06-12 07:09:43.816148: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2024-06-12 07:09:44.657802: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2251] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
2024-06-12 07:09:45.022935: E tensorflow/core/framework/dataset.cc:109] The Encode() method is not implemented for DatasetVariantWrapper objects.
2024-06-12 07:09:45.150111: E tensorflow/core/framework/variant.cc:100] Could not decode variant with type_name: "tensorflow::DatasetVariantWrapper".  Perhaps you forgot to register a decoder via REGISTER_UNARY_VARIANT_DECODE_FUNCTION?
2024-06-12 07:09:45.150452: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: DATA_LOSS: Unable to parse tensor from stored proto in file: /tmp/copycat/8099738303123461199/00000000.shard/00000000.snapshot, record 307. TensorProto: dtype: DT_VARIANT tensor_shape { } tensor_content: "#\n!tensorflow::DatasetVariantWrapper"
Traceback (most recent call last):
  File "/home/meiyang/src/ceprojects/magma/error.py", line 7, in <module>
    for x in ds2:
  File "/home/meiyang/miniconda3/envs/magma/lib/python3.10/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 809, in __next__
    return self._next_internal()
  File "/home/meiyang/miniconda3/envs/magma/lib/python3.10/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 772, in _next_internal
    ret = gen_dataset_ops.iterator_get_next(
  File "/home/meiyang/miniconda3/envs/magma/lib/python3.10/site-packages/tensorflow/python/ops/gen_dataset_ops.py", line 3086, in iterator_get_next
    _ops.raise_from_not_ok_status(e, name)
  File "/home/meiyang/miniconda3/envs/magma/lib/python3.10/site-packages/tensorflow/python/framework/ops.py", line 5983, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.DataLossError: {{function_node __wrapped__IteratorGetNext_output_types_3_device_/job:localhost/replica:0/task:0/device:CPU:0}} Unable to parse tensor from stored proto in file: /tmp/copycat/8099738303123461199/00000000.shard/00000000.snapshot, record 307. TensorProto: dtype: DT_VARIANT tensor_shape { } tensor_content: "#\n!tensorflow::DatasetVariantWrapper" [Op:IteratorGetNext] name:
```
