Tensorflow inference error
Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61496">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61496">No</a>

### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

tf 2.8

### Custom code

Yes

### OS platform and distribution

linux centos

### Mobile device

_No response_

### Python version

3.8

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

My model can train normally, but there was an error when inference after the training was completed.The model structure code is as follows:

conv_1 = Conv2D(32,(1,5),(1,1),name='mode0_conv_1',padding='same')(input_1)
bn_1 = BatchNormalizetion(name='mode0_bn_1')(conv_1)
out_1 = PReLU(shared_axes=[1,2])(bn_1)
out_2 = tf.reshape(........)(out_1)
dp_1 = LSTM(......)(out_2)

dp_o1 = Dense(32,)(dp_1)
dp_o2 = PReLU(shared_axes=[1,2])(dp_o1)

ls_o1 = LSTM(................)(dp_o2)

dp_o3 = Dense(2,)(ls_o1)

The error is as follows:

tensorflow.pyrhon.framework.errors_impl.InvalidArgumentError:Graph execution error:
....

Node:'model/model0_bn_1/FuseBatchNormV3'
scale must have the same number of elements as the channels of x, got 32 and 2
         [[{node model/model0_bn_1/FuseBatchNormV3}]] [op:__inference_predict_function_3355]

May I ask what caused this and if it can be resolved?

### Standalone code to reproduce the issue

```shell
No
```


### Relevant log output

_No response_