Repository: tensorflow
Issue: RAM memory leak with tf.function when training multiple models in a loop
Labels: stat:awaiting tensorflower,type:bug,comp:tf.function,TF 2.15

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

tf 2.15

### Custom code

Yes

### OS platform and distribution

Ubuntu 22.04.3 LTS

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

n/a

### GPU model and memory

n/a

### Current behavior?

When I train multiple models in a loop, if I decorate the `train()` function with `@tf.function`, then the memory usage keeps on increasing after each loop iteration, even when I delete the model at the end of each loop and clear the TensorFlow graph/session.

The memory leak does not occur when `@tf.function` is removed. However, model training performance is significantly slower.

Colab notebook to reproduce issue:
https://colab.research.google.com/drive/1sJsGmcFeZVx6ImNbqnBsgF_LzIyPxXPW?usp=sharing

### Standalone code to reproduce the issue

```python
import gc
import os

import psutil
import tensorflow as tf


class MyModel:
    def __init__(self):
        self.dnn = tf.keras.Sequential([
            tf.keras.layers.Dense(256),
            tf.keras.layers.Dense(256),
            tf.keras.layers.Dense(256),
            tf.keras.layers.Dense(1),
        ])
        self.optimizer = tf.optimizers.Adam()
    
    @tf.function   # if we remove this @tf.function decorator, then there is no memory leak
    def train(self, X):
        with tf.GradientTape() as tape:
            loss = tf.reduce_sum(self.dnn(X))
        grads = tape.gradient(loss, self.dnn.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.dnn.trainable_variables))

process = psutil.Process(os.getpid())
rss = int(process.memory_info().rss / 1024 / 1024)  # in MB
print(f'rss: {rss} MB')

X = tf.ones((50, 80))
for i in range(50):
    model = MyModel()
    for _ in range(20):
        model.train(X)

    del model
    tf.keras.backend.clear_session()
    tf.compat.v1.reset_default_graph()
    gc.collect()

    new_rss = int(process.memory_info().rss / 1024 / 1024)
    if new_rss > rss:
        rss_increase = new_rss - rss
        rss = new_rss
        print(f'Iter {i:05d}, rss increase: {rss_increase} MB, rss: {rss} MB')
```


### Relevant log output

```shell
rss: 593 MB
Iter 00000, rss increase: 31 MB, rss: 624 MB
Iter 00001, rss increase: 6 MB, rss: 630 MB
Iter 00002, rss increase: 5 MB, rss: 635 MB
Iter 00003, rss increase: 5 MB, rss: 640 MB
Iter 00004, rss increase: 5 MB, rss: 645 MB
Iter 00005, rss increase: 5 MB, rss: 650 MB
Iter 00006, rss increase: 5 MB, rss: 655 MB
Iter 00007, rss increase: 5 MB, rss: 660 MB
Iter 00008, rss increase: 5 MB, rss: 665 MB
Iter 00009, rss increase: 5 MB, rss: 670 MB
Iter 00010, rss increase: 4 MB, rss: 674 MB
---  <omitting some rows for brevity>  ---
Iter 00040, rss increase: 5 MB, rss: 822 MB
Iter 00041, rss increase: 5 MB, rss: 827 MB
Iter 00042, rss increase: 5 MB, rss: 832 MB
Iter 00043, rss increase: 5 MB, rss: 837 MB
Iter 00044, rss increase: 5 MB, rss: 842 MB
Iter 00045, rss increase: 5 MB, rss: 847 MB
Iter 00046, rss increase: 5 MB, rss: 852 MB
Iter 00047, rss increase: 5 MB, rss: 857 MB
Iter 00048, rss increase: 5 MB, rss: 862 MB
Iter 00049, rss increase: 5 MB, rss: 867 MB
```