jit_compile keeps filling gpu memory to crash with simple dense and small data with official docker image
Hi @de-mh ,

I have tested the code on Colab and it works fine as per attached [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/2f376aa2f5ebe8af3098caa4ac723ef5/63209.ipynb#scrollTo=O-oucVZCfKlL).Please note that Colab has Tesla T4 16GB GPU. Maybe this is related to docker image as the crash is after few epochs. Could you please confirm TF docker image source you have used?

Thanks! Hi @SuryanarayanaY 

This is the image that I used. It is the latest official image
```tensorflow/tensorflow   latest-gpu-jupyter    a2cf87758fef   3 months ago    7.32GB```
### Issue type

Bug

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

Yes

### Source

binary

### TensorFlow version

v2.15.0-2-g0b15fdfcb3f 2.15.0

### Custom code

No

### OS platform and distribution

Linux Ubuntu 22 (Docker)

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

12.3

### GPU model and memory

_No response_

### Current behavior?

To narrow down the problem, I simplified everything as much as possible. I used official latest docker image to isolate as much as possible.
This is the GPU stats:
```
Thu Mar  7 17:16:10 2024       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.161.07             Driver Version: 535.161.07   CUDA Version: 12.3     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla M60                      Off | 00000000:00:1E.0 Off |                    0 |
| N/A   30C    P8              14W / 150W |      0MiB /  7680MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|  No running processes found                                                           |
+---------------------------------------------------------------------------------------+
```

 I'm running These few simple lines:
```
import tensorflow as tf

physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

with tf.device("CPU"):
    train_dataset = tf.data.Dataset.from_tensor_slices((tf.random.uniform((1000_000,400)), tf.random.uniform((1000_000,1)))).batch(4000)

model = tf.keras.Sequential(
    [
        tf.keras.Input(shape=(400,)),
        tf.keras.layers.Dense(400, activation='relu'),
        tf.keras.layers.Dense(400, activation='relu'),
        tf.keras.layers.Dense(400, activation='relu'),
        tf.keras.layers.Dense(400, activation='relu'),
        tf.keras.layers.Dense(1, activation='sigmoid'),
    ]
)
model.compile(
    optimizer=tf.keras.optimizers.Adam(), 
    loss=tf.keras.losses.BinaryCrossentropy(from_logits=False), 
    jit_compile=True,
)

model.fit(train_dataset, 
            epochs=50, 
            verbose=True,
)
```

And the code crashes after 15 epochs because memory is filled. Here is the notebook for full errors detail:
[Break_XLA.ipynb.zip](https://github.com/tensorflow/tensorflow/files/14527818/Break_XLA.1.ipynb.zip)


### Standalone code to reproduce the issue

```shell
import tensorflow as tf

physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

with tf.device("CPU"):
    train_dataset = tf.data.Dataset.from_tensor_slices((tf.random.uniform((1000_000,400)), tf.random.uniform((1000_000,1)))).batch(4000)

model = tf.keras.Sequential(
    [
        tf.keras.Input(shape=(400,)),
        tf.keras.layers.Dense(400, activation='relu'),
        tf.keras.layers.Dense(400, activation='relu'),
        tf.keras.layers.Dense(400, activation='relu'),
        tf.keras.layers.Dense(400, activation='relu'),
        tf.keras.layers.Dense(1, activation='sigmoid'),
    ]
)
model.compile(
    optimizer=tf.keras.optimizers.Adam(), 
    loss=tf.keras.losses.BinaryCrossentropy(from_logits=False), 
    jit_compile=True,
)

model.fit(train_dataset, 
            epochs=50, 
            verbose=True,
)
```



### Relevant log output

```shell
Epoch 1/50
2024-03-07 17:28:52.944900: I external/local_xla/xla/service/service.cc:168] XLA service 0x7fd9180c3f10 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2024-03-07 17:28:52.944951: I external/local_xla/xla/service/service.cc:176]   StreamExecutor device (0): Tesla M60, Compute Capability 5.2
2024-03-07 17:28:52.980936: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2024-03-07 17:28:53.293838: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:454] Loaded cuDNN version 8906
 13/250 [>.............................] - ETA: 3s - loss: 0.7202
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1709832534.134273     948 device_compiler.h:186] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
250/250 [==============================] - 5s 11ms/step - loss: 0.6946
Epoch 2/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6932
Epoch 3/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6932
Epoch 4/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6932
Epoch 5/50
250/250 [==============================] - 2s 10ms/step - loss: 0.6932
Epoch 6/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6931
Epoch 7/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6931
Epoch 8/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6931
Epoch 9/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6931
Epoch 10/50
250/250 [==============================] - 2s 10ms/step - loss: 0.6931
Epoch 11/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6931
Epoch 12/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6931
Epoch 13/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6931
Epoch 14/50
250/250 [==============================] - 3s 10ms/step - loss: 0.6931
Epoch 15/50
193/250 [======================>.......] - ETA: 0s - loss: 0.6931
2024-03-07 17:29:31.761561: W external/local_xla/xla/service/gpu/runtime/support.cc:58] Intercepted XLA runtime error:
INTERNAL: the requested functionality is not supported
2024-03-07 17:29:31.761643: W external/local_xla/xla/service/gpu/runtime/support.cc:58] Intercepted XLA runtime error:
INTERNAL: CaptureGpuGraph failed (the requested functionality is not supported; current tracing scope: custom-call): INTERNAL: Failed to end stream capture: CUDA_ERROR_STREAM_CAPTURE_INVALIDATED: operation failed due to a previous error during capture
---------------------------------------------------------------------------
InternalError                             Traceback (most recent call last)
Cell In[5], line 17
      1 model = tf.keras.Sequential(
      2     [
      3         tf.keras.Input(shape=(400,)),
   (...)
      9     ]
     10 )
     11 model.compile(
     12     optimizer=tf.keras.optimizers.Adam(), 
     13     loss=tf.keras.losses.BinaryCrossentropy(from_logits=False), 
     14     jit_compile=True,
     15 )
---> 17 model.fit(train_dataset, 
     18             epochs=50, 
     19             verbose=True,
     20 )

File /usr/local/lib/python3.11/dist-packages/keras/src/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File /usr/local/lib/python3.11/dist-packages/tensorflow/python/eager/execute.py:53, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     51 try:
     52   ctx.ensure_initialized()
---> 53   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     54                                       inputs, attrs, num_outputs)
     55 except core._NotOkStatusException as e:
     56   if name is not None:

InternalError: Graph execution error:

Detected at node StatefulPartitionedCall defined at (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main

  File "<frozen runpy>", line 88, in _run_code

  File "/usr/local/lib/python3.11/dist-packages/ipykernel_launcher.py", line 17, in <module>

  File "/usr/local/lib/python3.11/dist-packages/traitlets/config/application.py", line 1077, in launch_instance

  File "/usr/local/lib/python3.11/dist-packages/ipykernel/kernelapp.py", line 739, in start

  File "/usr/local/lib/python3.11/dist-packages/tornado/platform/asyncio.py", line 205, in start

  File "/usr/lib/python3.11/asyncio/base_events.py", line 604, in run_forever

  File "/usr/lib/python3.11/asyncio/base_events.py", line 1909, in _run_once

  File "/usr/lib/python3.11/asyncio/events.py", line 80, in _run

  File "/usr/local/lib/python3.11/dist-packages/ipykernel/kernelbase.py", line 529, in dispatch_queue

  File "/usr/local/lib/python3.11/dist-packages/ipykernel/kernelbase.py", line 518, in process_one

  File "/usr/local/lib/python3.11/dist-packages/ipykernel/kernelbase.py", line 424, in dispatch_shell

  File "/usr/local/lib/python3.11/dist-packages/ipykernel/kernelbase.py", line 766, in execute_request

  File "/usr/local/lib/python3.11/dist-packages/ipykernel/ipkernel.py", line 429, in do_execute

  File "/usr/local/lib/python3.11/dist-packages/ipykernel/zmqshell.py", line 549, in run_cell

  File "/usr/local/lib/python3.11/dist-packages/IPython/core/interactiveshell.py", line 3048, in run_cell

  File "/usr/local/lib/python3.11/dist-packages/IPython/core/interactiveshell.py", line 3103, in _run_cell

  File "/usr/local/lib/python3.11/dist-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner

  File "/usr/local/lib/python3.11/dist-packages/IPython/core/interactiveshell.py", line 3308, in run_cell_async

  File "/usr/local/lib/python3.11/dist-packages/IPython/core/interactiveshell.py", line 3490, in run_ast_nodes

  File "/usr/local/lib/python3.11/dist-packages/IPython/core/interactiveshell.py", line 3550, in run_code

  File "/tmp/ipykernel_879/3821501683.py", line 17, in <module>

  File "/usr/local/lib/python3.11/dist-packages/keras/src/utils/traceback_utils.py", line 65, in error_handler

  File "/usr/local/lib/python3.11/dist-packages/keras/src/engine/training.py", line 1807, in fit

  File "/usr/local/lib/python3.11/dist-packages/keras/src/engine/training.py", line 1401, in train_function

  File "/usr/local/lib/python3.11/dist-packages/keras/src/engine/training.py", line 1384, in step_function

Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.graph.launch' failed: CaptureGpuGraph failed (the requested functionality is not supported; current tracing scope: custom-call): INTERNAL: Failed to end stream capture: CUDA_ERROR_STREAM_CAPTURE_INVALIDATED: operation failed due to a previous error during capture; current profiling annotation: XlaModule:#hlo_module=a_inference_run_step_1288__.701,program_id=132#.
	 [[{{node StatefulPartitionedCall}}]] [Op:__inference_train_function_1361]
```
