snapshoting failure on dataset made from generator
@sachinprasadhs  I was able to replicate this issue on colab using TF v[2.12](https://colab.research.google.com/gist/sushreebarsa/6bbdf455ce7f0ee9a8c21451e72ad619/61313.ipynb) ,[2.13](https://colab.research.google.com/gist/sushreebarsa/5b2c607a0a6cd2be0d368bad7c384cc6/61313.ipynb) and [tf-nightly](https://colab.research.google.com/gist/sushreebarsa/dddddf65fbceffcc20c8a9c434dc9d37/61313.ipynb#scrollTo=boqcOlPBfwc7). Please find the attached gists. Thank you! Hi. I am a developer looking to get into open source. Have been using Tensorflow for a long time and also thinking of contributing to the repository. I had some questions to ask regarding this issue. 
- According to assesment is this a good issue to start on. 
- If yes, is there some guide on how I can test my changes on code for the same code, locally. 
### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

v2.12.0-rc1-12-g0db597d0d75

### Custom code

Yes

### OS platform and distribution

Linux 4e51bcd72cb8 5.15.109 (Colab)

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Simple snapshoting op with sharding fails when applied to a generator-based dataset.

### Standalone code to reproduce the issue

```shell
import numpy as np
import tensorflow as tf

np.random.seed(1234)

IMG_SHAPE = (224,224,3)

def gen_img(shape=IMG_SHAPE):
  while True:
    img = np.random.randint(0,256,size=IMG_SHAPE)
    lab = np.random.randint(0,10)
    yield (img,lab)

ds = tf.data.Dataset.from_generator(
      gen_img,
      output_signature=(
        tf.TensorSpec(shape=IMG_SHAPE, dtype=tf.int32),
        tf.TensorSpec(shape=(), dtype=tf.int32)
      )
)
ds = ds.take(int(1e3)).batch(32)
ds = ds.enumerate()
ds = ds.snapshot('./my_cached_dataset', shard_func = lambda i,x: i%10)
ds = ds.map(lambda i,x: x).repeat(2) # error disappears under 1 epoch !

for i,(img,lab) in enumerate(ds):
  pass
```


### Relevant log output

```shell
ResourceExhaustedError: {{function_node __wrapped__IteratorGetNext_output_types_2_device_/job:localhost/replica:0/task:0/device:CPU:0}} Output buffer(size: 262144 bytes) too small. Should be larger than 19267611 bytes. [Op:IteratorGetNext]
```
