Repository: tensorflow
Issue: Unnecessary memcopies between CPU and GPU when using tf.function
Labels: stat:awaiting tensorflower,type:bug,type:performance,TF 2.12

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

Bug

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

No

### Source

source

### Tensorflow Version

2.12.0

### Custom Code

Yes

### OS Platform and Distribution

Linux CentOS

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

11.8

### GPU model and memory

_No response_

### Current Behaviour?

We recognized that TensorFlow creates unnecessary copies between CPU and GPU. To reproduce, save the code below as `error.py` and execute it using `nvprof --print-gpu-trace --openacc-profiling off python3 error.py MODE`, with mode being 0, 1, 2 or 3.

Here what we have observed, I simplyfied the profiler output and commented inline:

## Mode 0 (generates input on GPU, runs model on GPU using tf.function):
```python
    Size Name
## RANDOM NUMBER GENERATION ##
      8B [CUDA memcpy HtoD]
1.0039KB [CUDA memset]
       - void tensorflow::functor::FillPhiloxRandomKernelLaunch...
      8B [CUDA memcpy HtoD]
      8B [CUDA memcpy DtoH]
      8B [CUDA memcpy HtoD]
      4B [CUDA memcpy HtoD]
      8B [CUDA memcpy HtoD]
## HERE THE COMPUTATION STARTS
18.375MB [CUDA memcpy DtoH]	## TF copies data to CPU
18.375MB [CUDA memcpy HtoD] ## 1. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [928]
      8B [CUDA memcpy DtoD]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 2. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [963]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 3. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [997]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 4. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1031]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 5. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1063]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 6. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1099]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 7. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1133]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 8. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1167]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 9. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1201]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 10. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1235]
       - void Eigen::internal::EigenMetaKernel...
## Merging of all results into a single Tensor and ten copies back to host
     30B [CUDA memcpy HtoD]
       - void tensorflow::functor::ColumnReduceMax16ColumnsKernel...
       - void tensorflow::functor::BlockReduceKernel...
      1B [CUDA memcpy DtoH]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
183.75MB [CUDA memcpy DtoH]
```

So we see that TF copies the data, that is already on the GPU to the CPU, and then in every iteration copies back to GPU, instead of just using the data that is already on the GPU.

## Mode 1 (generates input on CPU, runs model on GPU using tf.function):
```python
    Size Name
      8B [CUDA memcpy HtoD]
1.0039KB [CUDA memset]
      8B [CUDA memcpy HtoD]
      8B [CUDA memcpy DtoH]
      8B [CUDA memcpy HtoD]
      4B [CUDA memcpy HtoD]
      8B [CUDA memcpy HtoD]
## HERE THE COMPUTATION STARTS
18.375MB [CUDA memcpy HtoD] ## 1. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [911]
      8B [CUDA memcpy DtoD]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 2. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [946]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 3. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [980]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 4. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1012]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 5. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1046]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 6. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1080]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 7. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1116]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 8. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1150]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 9. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1184]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy HtoD] ## 10. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1218]
       - void Eigen::internal::EigenMetaKernel...
## Merging of all results into a single Tensor and ten copies back to host
     30B [CUDA memcpy HtoD]
       - void tensorflow::functor::ColumnReduceMax16ColumnsKernel...
       - void tensorflow::functor::BlockReduceKernel...
      1B [CUDA memcpy DtoH]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
18.375MB [CUDA memcpy DtoD]
183.75MB [CUDA memcpy DtoH]
```

Nearly identical to Mode 0, but here it is actually expected that the data gets copied over from CPU to GPU in every iteration.

## Mode 2 (generates input on GPU, runs model on GPU using tf.function): 
```python
    Size Name
## RANDOM NUMBER GENERATION ##
      8B [CUDA memcpy HtoD]
1.0039KB [CUDA memset]
       - void tensorflow::functor::FillPhiloxRandomKernelLaunch...
## HERE THE COMPUTATION STARTS
18.375MB [CUDA memcpy DtoH] ## TF copies data to CPU
      4B [CUDA memcpy HtoD]
      8B [CUDA memcpy HtoD]
18.375MB [CUDA memcpy HtoD] ## 1. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [907]
      8B [CUDA memcpy DtoD]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy DtoH] ## 1. TF copies result to CPU
18.375MB [CUDA memcpy HtoD] ## 2. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [991]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy DtoH] ## 2. TF copies result to CPU
18.375MB [CUDA memcpy HtoD] ## 3. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1076]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy DtoH] ## 3. TF copies result to CPU
18.375MB [CUDA memcpy HtoD] ## 4. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1161]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy DtoH] ## 4. TF copies result to CPU
18.375MB [CUDA memcpy HtoD] ## 5. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1246]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy DtoH] ## 5. TF copies result to CPU
18.375MB [CUDA memcpy HtoD] ## 6. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1333]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy DtoH] ## 6. TF copies result to CPU
18.375MB [CUDA memcpy HtoD] ## 7. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1420]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy DtoH] ## 7. TF copies result to CPU
18.375MB [CUDA memcpy HtoD] ## 8. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1509]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy DtoH] ## 8. TF copies result to CPU
18.375MB [CUDA memcpy HtoD] ## 9. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1594]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy DtoH] ## 9. TF copies result to CPU
18.375MB [CUDA memcpy HtoD] ## 10. TF copies data to GPU
       - Mul_GPU_DT_FLOAT_DT_FLOAT_kernel [1679]
       - void Eigen::internal::EigenMetaKernel...
18.375MB [CUDA memcpy DtoH] ## 10. TF copies result to CPU
```

Similar to mode 0, but here the results get immediatly copied back to the host.

## Mode 3 (generates input on GPU, runs model on GPU using eager mode): 
```python
    Size Name
## RANDOM NUMBER GENERATION ##
      8B [CUDA memcpy HtoD]
1.0039KB [CUDA memset]
       - void tensorflow::functor::FillPhiloxRandomKernelLaunch...
## HERE THE COMPUTATION STARTS
       - AddV2_GPU_DT_FLOAT_DT_FLOAT_kernel [806]
       - AddV2_GPU_DT_FLOAT_DT_FLOAT_kernel [810]
       - AddV2_GPU_DT_FLOAT_DT_FLOAT_kernel [814]
       - AddV2_GPU_DT_FLOAT_DT_FLOAT_kernel [818]
       - AddV2_GPU_DT_FLOAT_DT_FLOAT_kernel [822]
       - AddV2_GPU_DT_FLOAT_DT_FLOAT_kernel [826]
       - AddV2_GPU_DT_FLOAT_DT_FLOAT_kernel [830]
       - AddV2_GPU_DT_FLOAT_DT_FLOAT_kernel [834]
       - AddV2_GPU_DT_FLOAT_DT_FLOAT_kernel [838]
       - AddV2_GPU_DT_FLOAT_DT_FLOAT_kernel [842]
```

In this case no uncessary memcopies occur, data is kept on GPU all the time.

## Summary
When we use `model.predict(...)` or `model.predict_on_batch()` (or any other of these Keras.Model functions that use `tf.function`), then the input data ALWAYS gets copied to the host first, and then in every iteration back to the GPU. This causes an significant performance penalty.

I have not been able to find any documentation about this behavior, if it is intended, or a way to prevent this to happen.

Here also the `tf.debugging.set_log_device_placement(True)` output. I think the line `input: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0` indicates that for whatever reason the `tf.function`'s input is expected to be on the CPU.

```python
2023-06-15 14:44:14.959681: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
resource_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.960017: I tensorflow/core/common_runtime/placer.cc:114] resource_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
VarHandleOp: (VarHandleOp): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.960041: I tensorflow/core/common_runtime/placer.cc:114] VarHandleOp: (VarHandleOp): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.960475: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op VarHandleOp in device /job:localhost/replica:0/task:0/device:CPU:0
resource: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.960895: I tensorflow/core/common_runtime/placer.cc:114] resource: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
value: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.960922: I tensorflow/core/common_runtime/placer.cc:114] value: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.960945: I tensorflow/core/common_runtime/placer.cc:114] AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.961429: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.963362: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
components_0: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.963837: I tensorflow/core/common_runtime/placer.cc:114] components_0: (_DeviceArg): /job:localhost/replica:0/task:0/device:CPU:0
TensorDataset: (TensorDataset): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.963872: I tensorflow/core/common_runtime/placer.cc:114] TensorDataset: (TensorDataset): /job:localhost/replica:0/task:0/device:CPU:0
handle_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.963889: I tensorflow/core/common_runtime/placer.cc:114] handle_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.964369: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op TensorDataset in device /job:localhost/replica:0/task:0/device:CPU:0
input__dataset: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.980981: I tensorflow/core/common_runtime/placer.cc:114] input__dataset: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
FlatMapDataset: (FlatMapDataset): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.981006: I tensorflow/core/common_runtime/placer.cc:114] FlatMapDataset: (FlatMapDataset): /job:localhost/replica:0/task:0/device:CPU:0
handle_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.981024: I tensorflow/core/common_runtime/placer.cc:114] handle_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.982115: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op FlatMapDataset in device /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.982841: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
input__dataset: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.983325: I tensorflow/core/common_runtime/placer.cc:114] input__dataset: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
buffer__size: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.983348: I tensorflow/core/common_runtime/placer.cc:114] buffer__size: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
PrefetchDataset: (PrefetchDataset): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.983387: I tensorflow/core/common_runtime/placer.cc:114] PrefetchDataset: (PrefetchDataset): /job:localhost/replica:0/task:0/device:CPU:0
handle_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.983403: I tensorflow/core/common_runtime/placer.cc:114] handle_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.984096: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op PrefetchDataset in device /job:localhost/replica:0/task:0/device:CPU:0
resource: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.984658: I tensorflow/core/common_runtime/placer.cc:114] resource: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
ReadVariableOp: (ReadVariableOp): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.984692: I tensorflow/core/common_runtime/placer.cc:114] ReadVariableOp: (ReadVariableOp): /job:localhost/replica:0/task:0/device:CPU:0
value_RetVal: (_DeviceRetval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.984710: I tensorflow/core/common_runtime/placer.cc:114] value_RetVal: (_DeviceRetval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.985212: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:CPU:0
input: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.985724: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
Identity: (Identity): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.985761: I tensorflow/core/common_runtime/placer.cc:114] Identity: (Identity): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.985778: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.986346: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.986722: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.986812: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.988278: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.988585: I tensorflow/core/common_runtime/placer.cc:114] resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.988624: I tensorflow/core/common_runtime/placer.cc:114] value: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.988643: I tensorflow/core/common_runtime/placer.cc:114] AssignVariableOp: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.989152: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op AssignVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.989472: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.989562: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:14.989661: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.989731: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
handle_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.990139: I tensorflow/core/common_runtime/placer.cc:114] handle_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
AnonymousIteratorV3: (AnonymousIteratorV3): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.990173: I tensorflow/core/common_runtime/placer.cc:114] AnonymousIteratorV3: (AnonymousIteratorV3): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.990613: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op AnonymousIteratorV3 in device /job:localhost/replica:0/task:0/device:CPU:0
dataset: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.990976: I tensorflow/core/common_runtime/placer.cc:114] dataset: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
iterator: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.990998: I tensorflow/core/common_runtime/placer.cc:114] iterator: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
MakeIterator: (MakeIterator): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.991016: I tensorflow/core/common_runtime/placer.cc:114] MakeIterator: (MakeIterator): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.991473: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op MakeIterator in device /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:14.992481: I tensorflow/core/common_runtime/executor.cc:1197] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_0' with dtype int32
         [[{{node Placeholder/_0}}]]
args_0: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.002874: I tensorflow/core/common_runtime/placer.cc:114] args_0: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
GeneratorDataset: (GeneratorDataset): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.002917: I tensorflow/core/common_runtime/placer.cc:114] GeneratorDataset: (GeneratorDataset): /job:localhost/replica:0/task:0/device:CPU:0
NoOp: (NoOp): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.002935: I tensorflow/core/common_runtime/placer.cc:114] NoOp: (NoOp): /job:localhost/replica:0/task:0/device:CPU:0
Identity: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.002966: I tensorflow/core/common_runtime/placer.cc:114] Identity: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
FakeSink0: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.002983: I tensorflow/core/common_runtime/placer.cc:114] FakeSink0: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.003007: I tensorflow/core/common_runtime/placer.cc:114] identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.007846: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op AnonymousIteratorV3 in device /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.007991: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op MakeIterator in device /job:localhost/replica:0/task:0/device:CPU:0
args_0: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.011294: I tensorflow/core/common_runtime/placer.cc:114] args_0: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
GeneratorDataset: (GeneratorDataset): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.011320: I tensorflow/core/common_runtime/placer.cc:114] GeneratorDataset: (GeneratorDataset): /job:localhost/replica:0/task:0/device:CPU:0
NoOp: (NoOp): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.011339: I tensorflow/core/common_runtime/placer.cc:114] NoOp: (NoOp): /job:localhost/replica:0/task:0/device:CPU:0
Identity: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.011355: I tensorflow/core/common_runtime/placer.cc:114] Identity: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
FakeSink0: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.011373: I tensorflow/core/common_runtime/placer.cc:114] FakeSink0: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.011389: I tensorflow/core/common_runtime/placer.cc:114] identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.018501: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.018640: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
iterator: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.046638: I tensorflow/core/common_runtime/placer.cc:114] iterator: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
assignaddvariableop_resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:15.046661: I tensorflow/core/common_runtime/placer.cc:114] assignaddvariableop_resource: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
IteratorGetNext: (IteratorGetNext): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.046678: I tensorflow/core/common_runtime/placer.cc:114] IteratorGetNext: (IteratorGetNext): /job:localhost/replica:0/task:0/device:CPU:0
model/tf.__operators__.add/AddV2: (AddV2): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:15.046695: I tensorflow/core/common_runtime/placer.cc:114] model/tf.__operators__.add/AddV2: (AddV2): /job:localhost/replica:0/task:0/device:GPU:0
AssignAddVariableOp: (AssignAddVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:15.046711: I tensorflow/core/common_runtime/placer.cc:114] AssignAddVariableOp: (AssignAddVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
NoOp: (NoOp): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:15.046726: I tensorflow/core/common_runtime/placer.cc:114] NoOp: (NoOp): /job:localhost/replica:0/task:0/device:GPU:0
Identity: (Identity): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:15.046759: I tensorflow/core/common_runtime/placer.cc:114] Identity: (Identity): /job:localhost/replica:0/task:0/device:GPU:0
identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:15.046783: I tensorflow/core/common_runtime/placer.cc:114] identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
Const: (Const): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:15.046798: I tensorflow/core/common_runtime/placer.cc:114] Const: (Const): /job:localhost/replica:0/task:0/device:GPU:0
2023-06-15 14:44:15.051534: I tensorflow/core/common_runtime/eager/execute.cc:1525] Executing op __inference_predict_function_68 in device /job:localhost/replica:0/task:0/device:GPU:0
args_0: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.054027: I tensorflow/core/common_runtime/placer.cc:114] args_0: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
NoOp: (NoOp): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.054141: I tensorflow/core/common_runtime/placer.cc:114] NoOp: (NoOp): /job:localhost/replica:0/task:0/device:CPU:0
Identity: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.054201: I tensorflow/core/common_runtime/placer.cc:114] Identity: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
FakeSink0: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.054272: I tensorflow/core/common_runtime/placer.cc:114] FakeSink0: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.054342: I tensorflow/core/common_runtime/placer.cc:114] identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
PyFunc: (PyFunc): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.054415: I tensorflow/core/common_runtime/placer.cc:114] PyFunc: (PyFunc): /job:localhost/replica:0/task:0/device:CPU:0
args_0: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.061278: I tensorflow/core/common_runtime/placer.cc:114] args_0: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
PyFunc: (PyFunc): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.061356: I tensorflow/core/common_runtime/placer.cc:114] PyFunc: (PyFunc): /job:localhost/replica:0/task:0/device:CPU:0
NoOp: (NoOp): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.061419: I tensorflow/core/common_runtime/placer.cc:114] NoOp: (NoOp): /job:localhost/replica:0/task:0/device:CPU:0
Identity: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.061476: I tensorflow/core/common_runtime/placer.cc:114] Identity: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
FakeSink0: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.061556: I tensorflow/core/common_runtime/placer.cc:114] FakeSink0: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.061612: I tensorflow/core/common_runtime/placer.cc:114] identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
args_0: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.067674: I tensorflow/core/common_runtime/placer.cc:114] args_0: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
PyFunc: (PyFunc): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.067767: I tensorflow/core/common_runtime/placer.cc:114] PyFunc: (PyFunc): /job:localhost/replica:0/task:0/device:CPU:0
NoOp: (NoOp): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.067834: I tensorflow/core/common_runtime/placer.cc:114] NoOp: (NoOp): /job:localhost/replica:0/task:0/device:CPU:0
Identity: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.067905: I tensorflow/core/common_runtime/placer.cc:114] Identity: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
FakeSink0: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.067969: I tensorflow/core/common_runtime/placer.cc:114] FakeSink0: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
2023-06-15 14:44:15.068020: I tensorflow/core/common_runtime/placer.cc:114] identity_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
```


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import sys

assert len(sys.argv) == 2, "needs to be run as `python3 error.py MODE`"

mode = int(sys.argv[1])

inp = tf.keras.Input((3, 224, 224))
out = inp + inp
model = tf.keras.Model(inp, out)

class Sequence(tf.keras.utils.Sequence):
        def __init__(self, x):          self.x = x
        def __len__(self):              return 10
        def __getitem__(self, idx):     return self.x

with tf.device('/GPU:0' if mode != 1 else '/CPU:0'):
        data = tf.random.uniform((32, 3, 224, 224))

with tf.device('/GPU:0'):
        if mode == 2:
                for _ in range(10):
                        model.predict_on_batch(data)
        elif mode == 3:
                for _ in range(10):
                        model(data)
        else:
                seq = Sequence(data)
                model.predict(seq)
```


### Relevant log output

```shell
see above
```
</details>