ValueError: as_list() is not defined on an unknown TensorShape. during training
Hi @guillaumelorre28 ,

I have reproduced the reported issue.Attached [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/580bb91c4f5cd5e76841016769ffdc61/68217.ipynb) for same.
### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

2.17.0-dev20240517

### Custom code

Yes

### OS platform and distribution

macOS Sonoma 

### Mobile device

_No response_

### Python version

3.10.13

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

I get the following error when training a transformer for translation:

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[9], [line 6](vscode-notebook-cell:?execution_count=9&line=6)
      [1](vscode-notebook-cell:?execution_count=9&line=1) model.compile(optimizer='adam', 
      [2](vscode-notebook-cell:?execution_count=9&line=2)     loss='sparse_categorical_crossentropy', 
      [3](vscode-notebook-cell:?execution_count=9&line=3)     metrics=['accuracy'],
      [4](vscode-notebook-cell:?execution_count=9&line=4)     run_eagerly=False
      [5](vscode-notebook-cell:?execution_count=9&line=5)     )
----> [6](vscode-notebook-cell:?execution_count=9&line=6) model.fit(dataset_train, epochs=10)

File ~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:122, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    [119](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:119)     filtered_tb = _process_traceback_frames(e.__traceback__)
    [120](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:120)     # To get the full stack trace, call:
    [121](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:121)     # `keras.config.disable_traceback_filtering()`
--> [122](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:122)     raise e.with_traceback(filtered_tb) from None
    [123](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:123) finally:
    [124](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py:124)     del filtered_tb

File ~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/optree/ops.py:594, in tree_map(func, tree, is_leaf, none_is_leaf, namespace, *rests)
    [592](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/optree/ops.py:592) leaves, treespec = _C.flatten(tree, is_leaf, none_is_leaf, namespace)
    [593](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/optree/ops.py:593) flat_args = [leaves] + [treespec.flatten_up_to(r) for r in rests]
--> [594](https://file+.vscode-resource.vscode-cdn.net/Users/glo/Desktop/deep_learning_training-main/~/opt/anaconda3/envs/tensorflow_latest/lib/python3.10/site-packages/optree/ops.py:594) return treespec.unflatten(map(func, *flat_args))

ValueError: as_list() is not defined on an unknown TensorShape.
```

### Standalone code to reproduce the issue

```shell
Here is a colab link to the code:
https://colab.research.google.com/drive/1BQ4lhaZPP5XGb_IUe-rVacMiFPWCfm9Y?usp=sharing

Running eagerly works fine but an error occurs in graph mode.
```


### Relevant log output

_No response_