customized model combined with keras model inputshape error
@MrWangg1992 Could you please provide the entire code? It appears that, during the model training, the `MobilenetV3small` layer receives an input whose last axis has singular dimensionality, while this layer expects input with a 3-dimensional last axis. There might be an issue with how the `CoordinateChannel` layer adds spatial information.
 > @MrWangg1992 Could you please provide the entire code? It appears that, during the model training, the `MobilenetV3small` layer receives an input whose last axis has singular dimensionality, while this layer expects input with a 3-dimensional last axis. There might be an issue with how the `CoordinateChannel` layer adds spatial information.

@aditya02shah 
Hi, from the error message, it said something regarding the dimensions. My original input image is a (B,224,224,1) then i send it to a coordinate channel change it to (B,224,224,3)  then send it to `MobilenetV3small`. This works well when in tf 2.3/2.4. but after I updated to tf 2.14/2.15 it threw up this error. As I mentioned, if I directly run model(input), it works fine, but after `model.compile()` ,`model.fit()`. it will give that error message.
I think i already provided the full code regarding how the model was built. the compile and fit part is not that necessary,  andy random `loss` `dataloader` is fine.

```
import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.regularizers import l2
import tensorflow as tf
from tensorflow.keras import backend as K
import numpy as np
import tensorflow.keras.layers as layers
import math


class CoordinateChannel(Layer):

    def __init__(self, rank,
                 use_radius=False,
                 data_format=None,
                 **kwargs):
        super(CoordinateChannel, self).__init__(**kwargs)

        if data_format not in [None, 'channels_first', 'channels_last']:
            raise ValueError('`data_format` must be either "channels_last", "channels_first" '
                             'or None.')

        self.rank = rank
        self.use_radius = use_radius
        self.data_format = K.image_data_format() if data_format is None else data_format
        self.axis = 1 if K.image_data_format() == 'channels_first' else -1

        self.input_spec = InputSpec(min_ndim=2)
        self.supports_masking = True

    def build(self, input_shape):
        assert len(input_shape) >= 2
        input_dim = input_shape[self.axis]

        self.input_spec = InputSpec(min_ndim=self.rank + 2,
                                    axes={self.axis: input_dim})
        self.built = True

    def call(self, inputs, training=None, mask=None):
        input_shape = K.shape(inputs)

        if self.rank == 1:
            input_shape = [input_shape[i] for i in range(3)]
            batch_shape, dim, channels = input_shape

            xx_range = K.tile(K.expand_dims(K.arange(0, dim, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            xx_range = K.expand_dims(xx_range, axis=-1)

            xx_channels = K.cast(xx_range, K.floatx())
            xx_channels = xx_channels / K.cast(dim - 1, K.floatx())
            xx_channels = (xx_channels * 2) - 1.

            outputs = K.concatenate([inputs, xx_channels], axis=-1)

        if self.rank == 2:
            if self.data_format == 'channels_first':
                inputs = K.permute_dimensions(inputs, [0, 2, 3, 1])
                input_shape = K.shape(inputs)

            input_shape = [input_shape[i] for i in range(4)]
            batch_shape, dim1, dim2, channels = input_shape

            xx_ones = tf.ones(K.stack([batch_shape, dim2]), dtype='float32')
            xx_ones = K.expand_dims(xx_ones, axis=-1)

            xx_range = K.tile(K.expand_dims(K.arange(0, dim1, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            xx_range = K.expand_dims(xx_range, axis=1)
            xx_channels = K.batch_dot(xx_ones, xx_range, axes=[2, 1])
            xx_channels = K.expand_dims(xx_channels, axis=-1)
            xx_channels = K.permute_dimensions(xx_channels, [0, 2, 1, 3])

            yy_ones = tf.ones(K.stack([batch_shape, dim1]), dtype='float32')
            yy_ones = K.expand_dims(yy_ones, axis=1)

            yy_range = K.tile(K.expand_dims(K.arange(0, dim2, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            yy_range = K.expand_dims(yy_range, axis=-1)

            yy_channels = K.batch_dot(yy_range, yy_ones, axes=[2, 1])
            yy_channels = K.expand_dims(yy_channels, axis=-1)
            yy_channels = K.permute_dimensions(yy_channels, [0, 2, 1, 3])

            xx_channels = K.cast(xx_channels, K.floatx())
            xx_channels = xx_channels / K.cast(dim1 - 1, K.floatx())
            xx_channels = (xx_channels * 2) - 1.

            yy_channels = K.cast(yy_channels, K.floatx())
            yy_channels = yy_channels / K.cast(dim2 - 1, K.floatx())
            yy_channels = (yy_channels * 2) - 1.

            # import pdb;pdb.set_trace()
            outputs = K.concatenate([inputs, xx_channels, yy_channels], axis=-1)
            # outputs = K.concatenate([inputs, tf.cast(xx_channels, dtype=tf.float16), tf.cast(yy_channels, dtype=tf.float16)], axis=-1)

            if self.use_radius:
                rr = K.sqrt(K.square(xx_channels - 0.5) +
                            K.square(yy_channels - 0.5))
                outputs = K.concatenate([outputs, rr], axis=-1)

            if self.data_format == 'channels_first':
                outputs = K.permute_dimensions(outputs, [0, 3, 1, 2])

        if self.rank == 3:
            if self.data_format == 'channels_first':
                inputs = K.permute_dimensions(inputs, [0, 2, 3, 4, 1])
                input_shape = K.shape(inputs)

            input_shape = [input_shape[i] for i in range(5)]
            batch_shape, dim1, dim2, dim3, channels = input_shape

            xx_ones = tf.ones(K.stack([batch_shape, dim3]), dtype='float32')
            xx_ones = K.expand_dims(xx_ones, axis=-1)

            xx_range = K.tile(K.expand_dims(K.arange(0, dim2, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            xx_range = K.expand_dims(xx_range, axis=1)

            xx_channels = K.batch_dot(xx_ones, xx_range, axes=[2, 1])
            xx_channels = K.expand_dims(xx_channels, axis=-1)
            xx_channels = K.permute_dimensions(xx_channels, [0, 2, 1, 3])

            xx_channels = K.expand_dims(xx_channels, axis=1)
            xx_channels = K.tile(xx_channels,
                                 [1, dim1, 1, 1, 1])

            yy_ones = tf.ones(K.stack([batch_shape, dim2]), dtype='float32')
            yy_ones = K.expand_dims(yy_ones, axis=1)

            yy_range = K.tile(K.expand_dims(K.arange(0, dim3, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            yy_range = K.expand_dims(yy_range, axis=-1)

            yy_channels = K.batch_dot(yy_range, yy_ones, axes=[2, 1])
            yy_channels = K.expand_dims(yy_channels, axis=-1)
            yy_channels = K.permute_dimensions(yy_channels, [0, 2, 1, 3])

            yy_channels = K.expand_dims(yy_channels, axis=1)
            yy_channels = K.tile(yy_channels,
                                 [1, dim1, 1, 1, 1])

            zz_range = K.tile(K.expand_dims(K.arange(0, dim1, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            zz_range = K.expand_dims(zz_range, axis=-1)
            zz_range = K.expand_dims(zz_range, axis=-1)

            zz_channels = K.tile(zz_range,
                                 [1, 1, dim2, dim3])
            zz_channels = K.expand_dims(zz_channels, axis=-1)

            xx_channels = K.cast(xx_channels, K.floatx())
            xx_channels = xx_channels / K.cast(dim2 - 1, K.floatx())
            xx_channels = xx_channels * 2 - 1.

            yy_channels = K.cast(yy_channels, K.floatx())
            yy_channels = yy_channels / K.cast(dim3 - 1, K.floatx())
            yy_channels = yy_channels * 2 - 1.

            zz_channels = K.cast(zz_channels, K.floatx())
            zz_channels = zz_channels / K.cast(dim1 - 1, K.floatx())
            zz_channels = zz_channels * 2 - 1.

            outputs = K.concatenate([inputs, zz_channels, xx_channels, yy_channels],
                                    axis=-1)

            if self.data_format == 'channels_first':
                outputs = K.permute_dimensions(outputs, [0, 4, 1, 2, 3])

        return outputs

    def compute_output_shape(self, input_shape):
        assert input_shape and len(input_shape) >= 2
        assert input_shape[self.axis]

        if self.use_radius and self.rank == 2:
            channel_count = 3
        else:
            channel_count = self.rank

        output_shape = list(input_shape)
        output_shape[self.axis] = input_shape[self.axis] + channel_count
        return tuple(output_shape)

    def get_config(self):
        config = {
            'rank': self.rank,
            'use_radius': self.use_radius,
            'data_format': self.data_format
        }
        base_config = super(CoordinateChannel, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

def binary_focal_loss(gamma=2., alpha=.25, amplifier=1.):
    def binary_focal_loss_fixed(y_true, y_pred):
        pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
        pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))

        epsilon = K.epsilon()
        # clip to prevent NaN's and Inf's
        pt_1 = K.clip(pt_1, epsilon, 1. - epsilon)
        pt_0 = K.clip(pt_0, epsilon, 1. - epsilon)
        return (-K.mean(alpha * K.pow(1. - pt_1, gamma) * K.log(pt_1), keepdims=True) \
                - K.mean((1 - alpha) * K.pow(pt_0, gamma) * K.log(1. - pt_0), keepdims=True)) * amplifier

    return binary_focal_loss_fixed

backbone = tf.keras.applications.MobileNetV3Small(
    input_shape=(224,224,3),
    alpha=1.0,
    minimalistic=False,
    include_top=True,
    weights='imagenet',
    input_tensor=None,
    classes=1000,
    pooling=None,
    dropout_rate=0.2,
    classifier_activation='softmax',
    include_preprocessing=True
)


inputs = Input(shape=[224, 224] + [1])
x = CoordinateChannel(2)(inputs)
backbone_out = backbone(x)

model = Model(inputs=inputs outputs=backbone_out)

lr = 0.001
opt = Adam(lr)
model.compile(optimizer=opt,loss=binary_focal_loss(gamma=5., alpha=0.5),metrics=[binary_focal_loss(gamma=5., alpha=0.5)])

model.fit(
    train_dataset,
    steps_per_epoch=train_steps,
    validation_data=val_dataset,
    validation_steps=val_steps,
    epochs=5,
    initial_epoch=0,
    workers=0
) # this part will throw error



ip = tf.random.normal((1, 224, 224, 1))
model(IP) # this works well 

``` @MrWangg1992, the code appears to be functioning with a random dataloader and loss that I tested. I recommend examining your specific dataloader implementation and how it interacts with the `CoordinateChannel` Layer. The issue might be related to the data input or preprocessing. Hope this helps!
[Gist](https://colab.research.google.com/drive/1LLrHA4b9JTvq3KJ3MhnbKB2Zs0zo91xi#scrollTo=UKurgsCWvkjj)
 > @MrWangg1992, the code appears to be functioning with a random dataloader and loss that I tested. I recommend examining your specific dataloader implementation and how it interacts with the `CoordinateChannel` Layer. The issue might be related to the data input or preprocessing. Hope this helps! [Gist](https://colab.research.google.com/drive/1LLrHA4b9JTvq3KJ3MhnbKB2Zs0zo91xi#scrollTo=UKurgsCWvkjj)

Hi @aditya02shah 
a quick question does, tf support a800 gpu? i think this is the root cause @MrWangg1992 To verify TensorFlow support for the A800 GPU, check out the documentation at these links: <br/>1.[NVIDIA CUDA GPUs](https://developer.nvidia.com/cuda-gpus) <br/> 2.[Tensorflow GPU Hardware Requirements](https://www.tensorflow.org/install/pip#hardware_requirements) > @MrWangg1992 To verify TensorFlow support for the A800 GPU, check out the documentation at these links: 1.[NVIDIA CUDA GPUs](https://developer.nvidia.com/cuda-gpus) 2.[Tensorflow GPU Hardware Requirements](https://www.tensorflow.org/install/pip#hardware_requirements)

@aditya02shah 
i didn't see a800 listed, would you mind check it on your end? i checked on my end and it doesn't work
![image](https://github.com/tensorflow/tensorflow/assets/31319863/5c69133b-eb5b-433e-9c82-c9e810b71bfc)
 @aditya02shah 
I also tried something else not working as well.
![image](https://github.com/tensorflow/tensorflow/assets/31319863/7593de11-5474-44c1-ae15-f82baea173d1)
 @MrWangg1992 Could you please refer to this [TF Forum Cuda Issue](https://discuss.tensorflow.org/t/tensorflow-install-error/20491) to see if it resolves your TensorFlow CUDA issue. Hi @MrWangg1992 ,

Could you please confirm did you installed `tensorflow[and-cuda]` package ? This will install all the cuda and cudnn packages required for GPU setup.For earlier versions TF<=2.13 it is manual and now its taken care by TF itself.

I would recommend you to try with new environment by installing `tensorflow[and-cuda]` and check whether GPU is recognisable.

Is the input_shape error resolved yet ?

Please submit a minimal code snippet for reproduction. For support type issue please post at [tensorflow-forum](https://discuss.tensorflow.org/). > Hi @MrWangg1992 ,
> 
> Could you please confirm did you installed `tensorflow[and-cuda]` package ? This will install all the cuda and cudnn packages required for GPU setup.For earlier versions TF<=2.13 it is manual and now its taken care by TF itself.
> 
> I would recommend you to try with new environment by installing `tensorflow[and-cuda]` and check whether GPU is recognisable.
> 
> Is the input_shape error resolved yet ?
> 
> Please submit a minimal code snippet for reproduction. For support type issue please post at [tensorflow-forum](https://discuss.tensorflow.org/).

Hi @SuryanarayanaY 
The gpu can be recognized now, but the shape error, didn't resolve, i used the same code, and everything works fine under tf 2.4/2.3, but I tried on tf 2.11/2.14/2.15, it throws different errors regarding the shape. I used the same processing methods above. i have no idea what changes inside, but it blocked me for a very long time. Hi @MrWangg1992 ,

Could you please submit minimal code snippet to reproduce the issue ?

Attached code is a large one and incomplete. If not a minimal code snippet please execute the code in google colab and submit the gist. This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. 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/62772">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/62772">No</a>

### Issue type

Bug

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

No

### Source

source

### TensorFlow version

2.15

### Custom code

Yes

### OS platform and distribution

linunx ubuntu 22.04

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

    model.fit(
  File "/usr/local/lib/python3.11/dist-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/tmp/__autograph_generated_filegq94afq9.py", line 15, in tf__train_function
    retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
    ^^^^^
ValueError: in user code:

    File "/usr/local/lib/python3.11/dist-packages/keras/src/engine/training.py", line 1401, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.11/dist-packages/keras/src/engine/training.py", line 1384, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.11/dist-packages/keras/src/engine/training.py", line 1373, in run_step  **
        outputs = model.train_step(data)
    File "/usr/local/lib/python3.11/dist-packages/keras/src/engine/training.py", line 1150, in train_step
        y_pred = self(x, training=True)
    File "/usr/local/lib/python3.11/dist-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/usr/local/lib/python3.11/dist-packages/keras/src/engine/input_spec.py", line 280, in assert_input_compatibility
        raise ValueError(

    ValueError: Exception encountered when calling layer 'MobilenetV3small' (type Functional).

    Input 0 of layer "Conv0" is incompatible with the layer: expected axis -1 of input shape to have value 3, but received input with shape (None, None, None, 1)

    Call arguments received by layer 'MobilenetV3small' (type Functional):
      • inputs=tf.Tensor(shape=(None, None, None, 1), dtype=float32)
      • training=True
      • mask=None

![image](https://github.com/tensorflow/tensorflow/assets/31319863/0c1899ff-1e74-4145-b378-142a7eb76305)

      
      
      

### Standalone code to reproduce the issue

```shell
class CoordinateChannel(Layer):
    """ Adds Coordinate Channels to the input tensor.
    # Arguments
        rank: An integer, the rank of the input data-uniform,
            e.g. "2" for 2D convolution.
        use_radius: Boolean flag to determine whether the
            radius coordinate should be added for 2D rank
            inputs or not.
        data_format: A string,
            one of `"channels_last"` or `"channels_first"`.
            The ordering of the dimensions in the inputs.
            `"channels_last"` corresponds to inputs with shape
            `(batch, ..., channels)` while `"channels_first"` corresponds to
            inputs with shape `(batch, channels, ...)`.
            It defaults to the `image_data_format` value found in your
            Keras config file at `~/.keras/keras.json`.
            If you never set it, then it will be "channels_last".
    # Input shape
        ND tensor with shape:
        `(samples, channels, *)`
        if `data_format` is `"channels_first"`
        or ND tensor with shape:
        `(samples, *, channels)`
        if `data_format` is `"channels_last"`.
    # Output shape
        ND tensor with shape:
        `(samples, channels + 2, *)`
        if `data_format` is `"channels_first"`
        or 5D tensor with shape:
        `(samples, *, channels + 2)`
        if `data_format` is `"channels_last"`.
    # References:
        - [An Intriguing Failing of Convolutional Neural Networks and the CoordConv Solution](https://arxiv.org/abs/1807.03247)
    """

    def __init__(self, rank,
                 use_radius=False,
                 data_format=None,
                 **kwargs):
        super(CoordinateChannel, self).__init__(**kwargs)

        if data_format not in [None, 'channels_first', 'channels_last']:
            raise ValueError('`data_format` must be either "channels_last", "channels_first" '
                             'or None.')

        self.rank = rank
        self.use_radius = use_radius
        self.data_format = K.image_data_format() if data_format is None else data_format
        self.axis = 1 if K.image_data_format() == 'channels_first' else -1

        self.input_spec = InputSpec(min_ndim=2)
        self.supports_masking = True

    def build(self, input_shape):
        assert len(input_shape) >= 2
        input_dim = input_shape[self.axis]

        self.input_spec = InputSpec(min_ndim=self.rank + 2,
                                    axes={self.axis: input_dim})
        self.built = True

    def call(self, inputs, training=None, mask=None):
        input_shape = K.shape(inputs)

        if self.rank == 1:
            input_shape = [input_shape[i] for i in range(3)]
            batch_shape, dim, channels = input_shape

            xx_range = K.tile(K.expand_dims(K.arange(0, dim, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            xx_range = K.expand_dims(xx_range, axis=-1)

            xx_channels = K.cast(xx_range, K.floatx())
            xx_channels = xx_channels / K.cast(dim - 1, K.floatx())
            xx_channels = (xx_channels * 2) - 1.

            outputs = K.concatenate([inputs, xx_channels], axis=-1)

        if self.rank == 2:
            if self.data_format == 'channels_first':
                inputs = K.permute_dimensions(inputs, [0, 2, 3, 1])
                input_shape = K.shape(inputs)

            input_shape = [input_shape[i] for i in range(4)]
            batch_shape, dim1, dim2, channels = input_shape

            xx_ones = tf.ones(K.stack([batch_shape, dim2]), dtype='float32')
            xx_ones = K.expand_dims(xx_ones, axis=-1)

            xx_range = K.tile(K.expand_dims(K.arange(0, dim1, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            xx_range = K.expand_dims(xx_range, axis=1)
            xx_channels = K.batch_dot(xx_ones, xx_range, axes=[2, 1])
            xx_channels = K.expand_dims(xx_channels, axis=-1)
            xx_channels = K.permute_dimensions(xx_channels, [0, 2, 1, 3])

            yy_ones = tf.ones(K.stack([batch_shape, dim1]), dtype='float32')
            yy_ones = K.expand_dims(yy_ones, axis=1)

            yy_range = K.tile(K.expand_dims(K.arange(0, dim2, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            yy_range = K.expand_dims(yy_range, axis=-1)

            yy_channels = K.batch_dot(yy_range, yy_ones, axes=[2, 1])
            yy_channels = K.expand_dims(yy_channels, axis=-1)
            yy_channels = K.permute_dimensions(yy_channels, [0, 2, 1, 3])

            xx_channels = K.cast(xx_channels, K.floatx())
            xx_channels = xx_channels / K.cast(dim1 - 1, K.floatx())
            xx_channels = (xx_channels * 2) - 1.

            yy_channels = K.cast(yy_channels, K.floatx())
            yy_channels = yy_channels / K.cast(dim2 - 1, K.floatx())
            yy_channels = (yy_channels * 2) - 1.

            # import pdb;pdb.set_trace()
            outputs = K.concatenate([inputs, xx_channels, yy_channels], axis=-1)
            # outputs = K.concatenate([inputs, tf.cast(xx_channels, dtype=tf.float16), tf.cast(yy_channels, dtype=tf.float16)], axis=-1)

            if self.use_radius:
                rr = K.sqrt(K.square(xx_channels - 0.5) +
                            K.square(yy_channels - 0.5))
                outputs = K.concatenate([outputs, rr], axis=-1)

            if self.data_format == 'channels_first':
                outputs = K.permute_dimensions(outputs, [0, 3, 1, 2])

        if self.rank == 3:
            if self.data_format == 'channels_first':
                inputs = K.permute_dimensions(inputs, [0, 2, 3, 4, 1])
                input_shape = K.shape(inputs)

            input_shape = [input_shape[i] for i in range(5)]
            batch_shape, dim1, dim2, dim3, channels = input_shape

            xx_ones = tf.ones(K.stack([batch_shape, dim3]), dtype='float32')
            xx_ones = K.expand_dims(xx_ones, axis=-1)

            xx_range = K.tile(K.expand_dims(K.arange(0, dim2, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            xx_range = K.expand_dims(xx_range, axis=1)

            xx_channels = K.batch_dot(xx_ones, xx_range, axes=[2, 1])
            xx_channels = K.expand_dims(xx_channels, axis=-1)
            xx_channels = K.permute_dimensions(xx_channels, [0, 2, 1, 3])

            xx_channels = K.expand_dims(xx_channels, axis=1)
            xx_channels = K.tile(xx_channels,
                                 [1, dim1, 1, 1, 1])

            yy_ones = tf.ones(K.stack([batch_shape, dim2]), dtype='float32')
            yy_ones = K.expand_dims(yy_ones, axis=1)

            yy_range = K.tile(K.expand_dims(K.arange(0, dim3, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            yy_range = K.expand_dims(yy_range, axis=-1)

            yy_channels = K.batch_dot(yy_range, yy_ones, axes=[2, 1])
            yy_channels = K.expand_dims(yy_channels, axis=-1)
            yy_channels = K.permute_dimensions(yy_channels, [0, 2, 1, 3])

            yy_channels = K.expand_dims(yy_channels, axis=1)
            yy_channels = K.tile(yy_channels,
                                 [1, dim1, 1, 1, 1])

            zz_range = K.tile(K.expand_dims(K.arange(0, dim1, dtype='float32'), axis=0),
                              K.stack([batch_shape, 1]))
            zz_range = K.expand_dims(zz_range, axis=-1)
            zz_range = K.expand_dims(zz_range, axis=-1)

            zz_channels = K.tile(zz_range,
                                 [1, 1, dim2, dim3])
            zz_channels = K.expand_dims(zz_channels, axis=-1)

            xx_channels = K.cast(xx_channels, K.floatx())
            xx_channels = xx_channels / K.cast(dim2 - 1, K.floatx())
            xx_channels = xx_channels * 2 - 1.

            yy_channels = K.cast(yy_channels, K.floatx())
            yy_channels = yy_channels / K.cast(dim3 - 1, K.floatx())
            yy_channels = yy_channels * 2 - 1.

            zz_channels = K.cast(zz_channels, K.floatx())
            zz_channels = zz_channels / K.cast(dim1 - 1, K.floatx())
            zz_channels = zz_channels * 2 - 1.

            outputs = K.concatenate([inputs, zz_channels, xx_channels, yy_channels],
                                    axis=-1)

            if self.data_format == 'channels_first':
                outputs = K.permute_dimensions(outputs, [0, 4, 1, 2, 3])

        return outputs

    def compute_output_shape(self, input_shape):
        assert input_shape and len(input_shape) >= 2
        assert input_shape[self.axis]

        if self.use_radius and self.rank == 2:
            channel_count = 3
        else:
            channel_count = self.rank

        output_shape = list(input_shape)
        output_shape[self.axis] = input_shape[self.axis] + channel_count
        return tuple(output_shape)

    def get_config(self):
        config = {
            'rank': self.rank,
            'use_radius': self.use_radius,
            'data_format': self.data_format
        }
        base_config = super(CoordinateChannel, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))



backbone = tf.keras.applications.MobileNetV3Small(
    input_shape=(224,224,3),
    alpha=1.0,
    minimalistic=False,
    include_top=True,
    weights='imagenet',
    input_tensor=None,
    classes=1000,
    pooling=None,
    dropout_rate=0.2,
    classifier_activation='softmax',
    include_preprocessing=True
)


inputs = Input(shape=[224, 224] + [1])
x = CoordinateChannel(2)(inputs)
backbone_out = backbone(x)

model = Model(inputs=inputs outputs=backbone_out)


ip = tf.random.normal((1, 224, 224, 1))
model(IP) # this works well 


but when i used for training, model. fir will throw the shape error
```


### Relevant log output

_No response_