Repository: tensorflow
Issue: Memory leak in custom training loop + tf.function 
Labels: stat:awaiting tensorflower,type:performance,comp:tf.function,TF 2.9

<em>Please make sure that this is an issue related to performance of TensorFlow.
As per our
[GitHub Policy](https://github.com/tensorflow/tensorflow/blob/master/ISSUES.md),
we only address code/doc bugs, performance issues, feature requests and
build/installation issues on GitHub. tag:performance_template</em>

**System information**

- Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18 (google colab)
- TensorFlow installed from (source or binary): pip
- TensorFlow version (use command below):  v2.5.0-0-ga4dfb8d1a71 2.5.0
- Python version: 3.7
- CUDA/cuDNN version: Build cuda_11.0_bu.TC445_37.28845127_0
- GPU model and memory: Tesla T4 - 15109MiB

You can collect some of this information using our environment capture
[script](https://github.com/tensorflow/tensorflow/tree/master/tools/tf_env_collect.sh)
You can also obtain the TensorFlow version with:
1. TF 1.0: `python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"`
2. TF 2.0: `python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"`

**Describe the current behavior**

I'm encountering a subtle memory leak, and unable to determine the source using tracemalloc. I run the code in google colab, which is meant to optimize hyperparameters for a custom ppo agent. Also, the speed at which the leak happens varies: sometimes it happens within 10-20 minutes of runtime / 5-10 iterations while on other times it may take up to an hour.

**Describe the expected behavior**

Memory usage should be constant over time, since there are no objects stored in memory intentionally as the training progresses.

**Standalone code to reproduce the issue**

Here's a colab [notebook](https://colab.research.google.com/drive/1mIO3Pd-HSskvIrcl3Bp3Evg_MzAEuwPc?usp=sharing)

**Other info / logs**

And here are the resulting [memory snapshots](https://drive.google.com/file/d/1RH89X5B5xsUY8a1RA4Gh9WmEIqXhs84R/view?usp=sharing) at 15 consecutive iterations up to the crash, which do not show any particular red flags. Besides, by summing up the `size` column of the most recent snapshot, it totals 967249710 bytes ~= 1GB, which is weird because the available memory on colab ~= 12GB. 

Here's a log after crashing:


```
   Timestamp                  Level    Message

Jul 14, 2021, 11:07:41 AM | WARNING | WARNING:root:kernel a3ecdd9e-1765-4f67-af0e-60c198c896ec restarted
Jul 14, 2021, 11:07:41 AM | INFO | KernelRestarter: restarting kernel (1/5), keep random ports
Jul 14, 2021, 10:22:16 AM | WARNING | 2021-07-14 08:22:16.518184: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcublasLt.so.11
Jul 14, 2021, 10:22:13 AM | WARNING | 2021-07-14 08:22:13.939377: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcublas.so.11
Jul 14, 2021, 10:21:50 AM | WARNING | 2021-07-14 08:21:50.140299: I tensorflow/stream_executor/cuda/cuda_dnn.cc:359] Loaded cuDNN version 8004
Jul 14, 2021, 10:21:47 AM | WARNING | 2021-07-14 08:21:47.841310: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudnn.so.8
Jul 14, 2021, 10:21:44 AM | WARNING | 2021-07-14 08:21:44.642335: I tensorflow/core/platform/profile_utils/cpu_utils.cc:114] CPU Frequency: 2199995000 Hz
Jul 14, 2021, 10:21:44 AM | WARNING | 2021-07-14 08:21:44.559306: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:176] None of the MLIR Optimization Passes are enabled (registered 2)
Jul 14, 2021, 10:21:30 AM | WARNING | 2021-07-14 08:21:30.355834: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1418] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 13837 MB memory) -> physical GPU (device: 0, name: Tesla T4, pci bus id: 0000:00:04.0, compute capability: 7.5)
Jul 14, 2021, 10:21:30 AM | WARNING | 2021-07-14 08:21:30.355771: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:39] Overriding allow_growth setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.
Jul 14, 2021, 10:21:30 AM | WARNING | 2021-07-14 08:21:30.354979: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
Jul 14, 2021, 10:21:30 AM | WARNING | 2021-07-14 08:21:30.354000: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
Jul 14, 2021, 10:21:30 AM | WARNING | 2021-07-14 08:21:30.352104: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
Jul 14, 2021, 10:21:30 AM | WARNING | 2021-07-14 08:21:30.351844: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1277] 0: N
Jul 14, 2021, 10:21:30 AM | WARNING | 2021-07-14 08:21:30.351828: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1264] 0
Jul 14, 2021, 10:21:30 AM | WARNING | 2021-07-14 08:21:30.351772: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1258] Device interconnect StreamExecutor with strength 1 edge matrix:
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.206565: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.203344: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1871] Adding visible gpu devices: 0
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.202526: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.201577: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
Jul 14, 2021, 10:21:25 AM | WARNING | coreClock: 1.59GHz coreCount: 40 deviceMemorySize: 14.75GiB deviceMemoryBandwidth: 298.08GiB/s
Jul 14, 2021, 10:21:25 AM | WARNING | pciBusID: 0000:00:04.0 name: Tesla T4 computeCapability: 7.5
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.201438: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1733] Found device 0 with properties:
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.200536: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.199581: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1871] Adding visible gpu devices: 0
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.195723: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.194598: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.194398: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudnn.so.8
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.190679: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcusparse.so.11
Jul 14, 2021, 10:21:25 AM | WARNING | 2021-07-14 08:21:25.144200: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcusolver.so.10
Jul 14, 2021, 10:21:24 AM | WARNING | 2021-07-14 08:21:24.851273: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcurand.so.10
Jul 14, 2021, 10:21:24 AM | WARNING | 2021-07-14 08:21:24.832903: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcufft.so.10
Jul 14, 2021, 10:21:24 AM | WARNING | 2021-07-14 08:21:24.660570: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcublasLt.so.11
Jul 14, 2021, 10:21:24 AM | WARNING | 2021-07-14 08:21:24.660379: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcublas.so.11
Jul 14, 2021, 10:21:24 AM | WARNING | 2021-07-14 08:21:24.533359: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
Jul 14, 2021, 10:21:24 AM | WARNING | coreClock: 1.59GHz coreCount: 40 deviceMemorySize: 14.75GiB deviceMemoryBandwidth: 298.08GiB/s
Jul 14, 2021, 10:21:24 AM | WARNING | pciBusID: 0000:00:04.0 name: Tesla T4 computeCapability: 7.5
Jul 14, 2021, 10:21:24 AM | WARNING | 2021-07-14 08:21:24.533253: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1733] Found device 0 with properties:
Jul 14, 2021, 10:21:24 AM | WARNING | 2021-07-14 08:21:24.532176: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
Jul 14, 2021, 10:21:24 AM | WARNING | 2021-07-14 08:21:24.461038: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcuda.so.1
Jul 14, 2021, 10:21:07 AM | WARNING | 2021-07-14 08:21:07.977292: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
Jul 14, 2021, 10:20:27 AM | INFO | Adapting to protocol v5.1 for kernel a3ecdd9e-1765-4f67-af0e-60c198c896ec
Jul 14, 2021, 10:20:26 AM | INFO | Kernel started: a3ecdd9e-1765-4f67-af0e-60c198c896ec
Jul 14, 2021, 10:20:20 AM | INFO | Use Control-C to stop this server and shut down all kernels (twice to skip confirmation).
Jul 14, 2021, 10:20:20 AM | INFO | http://172.28.0.12:9000/
Jul 14, 2021, 10:20:20 AM | INFO | The Jupyter Notebook is running at:
Jul 14, 2021, 10:20:20 AM | INFO | Use Control-C to stop this server and shut down all kernels (twice to skip confirmation).
Jul 14, 2021, 10:20:20 AM | INFO | 0 active kernels
Jul 14, 2021, 10:20:20 AM | INFO | http://172.28.0.2:9000/
Jul 14, 2021, 10:20:20 AM | INFO | The Jupyter Notebook is running at:
Jul 14, 2021, 10:20:20 AM | INFO | Serving notebooks from local directory: /
Jul 14, 2021, 10:20:20 AM | INFO | 0 active kernels
Jul 14, 2021, 10:20:20 AM | INFO | Serving notebooks from local directory: /
Jul 14, 2021, 10:20:20 AM | INFO | google.colab serverextension initialized.
Jul 14, 2021, 10:20:20 AM | INFO | google.colab serverextension initialized.
Jul 14, 2021, 10:20:20 AM | INFO | Writing notebook server cookie secret to /root/.local/share/jupyter/runtime/notebook_cookie_secret
Jul 14, 2021, 10:20:20 AM | INFO | Writing notebook server cookie secret to /root/.local/share/jupyter/runtime/notebook_cookie_secret
```