TF dataset generator memory leak?
@dhruvsreenivas Thanks for reporting issue.

Posting relevant issue threads [#58606](https://github.com/tensorflow/tensorflow/issues/58606), [#57690](https://github.com/tensorflow/tensorflow/issues/57690), [56624](https://github.com/tensorflow/tensorflow/issues/56624).

@sushreebarsa Could you please look into this. Thank you. @pjpratik thank you! I'll check these out.

If you guys need any help with error reproduction in Colab, let me know and I'll be glad to help. Alright, so as far as I can see, there is no definite solution posted -- I think some issues haven't necessarily been resolved yet, but they are shown to be legitimate (i.e. reproduced). What is the best course of action here? I'd like to know what to do as soon as possible. Thank you so much! @dhruvsreenivas Sorry for the late response!
Could you please share a simple standalone code to replicate this issue?
Thank you! Alright so the codebase is a bit big -- there are multiple files that generally contribute to the issue (i.e. parts from all these files go toward running the training loop, which will end in an OOM error), all generally organized in this one github [repo](https://github.com/dhruvsreenivas/byol-offline). I can give instructions on how to install everything and run the repo code if you want, but that might not be the best use of time.

I'm a bit busy now with some other work, so give me some time--I'll see if I can make a colab notebook showing everything over the weekend, and get back to you. How does that sound? Hi, @dhruvsreenivas 

Apologize for the delay and as you mentioned in your previous comment, you'll share one Google Colab notebook to show us TF dataset generator memory leak issue so May I know have you prepared that Google Colab notebook now ? If yes please share that notebook with us, we'll try to replicate the same issue from our end. Thank you! Hey everyone,

Sorry I've been super busy with some school work over the past few weeks -- I'll try to make the Colab by early next week, and will send upon completion. Thanks so much!

Best,
Dhruv Sreenivas Alright, I think it might not be a TensorFlow problem as much as it could just be I'm loading in a LOT of data at once, so memory might be at a premium -- I've tried this with a couple other experiments and it seems that this is what is going on. Sorry for the inconvenience!

For example, I tried something simple where I would delete the batch once I finished updating with it, and that actually worked well with limited memory. Hi, @dhruvsreenivas 

No problem at all, Good to hear that, You're able to resolve your issue, Thank you for sharing the detailed information about root cause for your issue and If you're having training data simply cannot fit in memory then I would suggest you to have look at this article [tf.data.Dataset generators with parallelization](https://medium.com/@acordier/tf-data-dataset-generators-with-parallelization-the-easy-way-b5c5f7d2a18) which may help you to solve your memory related issue.

If you need further assistance please let us know ? or Could you please confirm if this issue is resolved for you ? Please feel free to close the issue if it is resolved ? Thank you! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/59587">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/59587">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Performance

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

tf 2.9

### Custom Code

Yes

### OS Platform and Distribution

Linux Ubuntu 20.04

### Mobile device

_No response_

### Python version

3.9

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

NVIDIA RTX 2080ti (1 GPU used)

### Current Behaviour?

```shell
Hi everyone,

Hope you are doing well! I have an issue where potentially a memory leak occurs when I create a TensorFlow sampling dataset via tf.data.Dataset.from_generator. I am currently working on some model-based RL ideas, focusing on using exploration inside a Dreamer VAE model for best results in visual continuous control. The model is written in JAX/Haiku, and I decided to try to use a TensorFlow dataset pipeline to load in the data and use it to train the model.

The data is organized in trajectory format in a directory (400 trajectories, stored as .npz files, are in a directory for use). I've written code for the sampling procedure (see line 86 for the visual sequence buffer impl, line 349 for dataset/loader creation at https://github.com/dhruvsreenivas/byol-offline/blob/main/memory/replay_buffer.py for details), which I've done some preliminary tests on in line 170 of this file: https://github.com/dhruvsreenivas/byol-offline/blob/main/testing.py).

The behavior I'm seeing is that the model can effectively train for around 50 epochs with 50 GB of RAM requested, but jobs ultimately die due to memory leakage after that. I want to train the model for 1000 epochs, which makes it seem like an obscene amount of memory is required, which I don't have available to me.

I have a suspicion that a memory leak is occurring somewhere, which is odd given that I believe that memory should be dynamically allocated based off of the size of the batch requested (I am sampling a batch of 50 subtrajectories, each of length 10 for the initial experiments). I'm wondering how to fix this problem (the sampling is required, so the from_generator part should stay the same I think).
```


### Standalone code to reproduce the issue

```shell
The code to reproduce the experiments is here:
- Replay buffer code: https://github.com/dhruvsreenivas/byol-offline/blob/main/memory/replay_buffer.py -- line 86, 349
- Testing file: https://github.com/dhruvsreenivas/byol-offline/blob/main/testing.py -- line 170

For data I'm using to run the experiments, please see https://github.com/conglu1997/v-d4rl for how to collect and organize the data. I'm using cheetah_run 64x64px, medium expert data directory, seen here: https://drive.google.com/drive/folders/1zbAqR0gYBNG2W-F5_ItweCGg6_YHpZ7o.
```


### Relevant log output

```shell
I only really have the job's code, which basically says that the out-of-memory handler was used because the job ran out of memory. Used 2 CPUs and 1 GPU, with 50 GB RAM.
```
</details>