Tensorflow error: Input to reshape is a tensor with 4 values, but the requested shape has 1

Question

I'm trying to calculate the loss of my implementation of Mobilenet model. The loss function is as below:

def mobilenet_loss(y_true, y_pred):

    '''
    y_pred will be an tensor of (None, 21, 21, 3) where the last axis corresponds to the classes from class_list

    y_true will be an integer 0, 1 or 2 that would indicate which class the true value belongs to

    Since y_pred is the output of a Softmax function, the sum of elements in the last dim will be 1.0
    Using y_true, we will construct a tensor y_true_array with ones in the index of y_true and zeros in others
    We will be trying to minimize sum(y_true_array - y_pred)

    '''
    
    diagnostics = True
    if diagnostics: print("ytrue =", y_true, " y_pred shape =", y_pred.shape)
    y_true_array= tf.zeros_like(y_pred, dtype=tf.float32, name="loss_tru_array")
    one_class_shape = y_pred.shape[:-1] + [1]
    if diagnostics: print("one_class_shape =", one_class_shape)
    ones =tf.ones(one_class_shape, dtype=tf.float32, name="loss_ones")
    zeros = tf.zeros(one_class_shape, dtype=tf.float32, name="loss_ones")
    if (tf.math.equal(y_true, tf.constant(0, dtype=tf.int64))):
        y_true_array = tf.concat([ones, zeros, zeros], axis=-1)
    elif (tf.math.equal(y_true, tf.constant(1, dtype=tf.int64))):
        y_true_array = tf.concat([zeros, ones, zeros, ], axis=-1)
    elif (tf.math.equal(y_true, tf.constant(2 , dtype=tf.int64))):
        y_true_array = tf.concat([zeros, zeros, ones], axis=-1)
    if diagnostics:
        print("after y_true_array =", y_true_array.shape)

    error = tf.math.subtract(y_true_array, y_pred, name="error")
    if diagnostics: print("error shape = ", error.shape)
    loss = tf.math.reduce_sum(error, name='final_loss')
    if diagnostics: print("loss shape = ", loss.shape, loss)
    return loss

I run unit test of this function using the code below to get the output as shown following that

y_true = 1
y_pred = tf.random.normal(shape = [1, 21,21,3], seed = 7)
out = mobilenet_loss(y_true, y_pred)
out.shape

Output of the unit test:

ytrue = 1  y_pred shape = (1, 21, 21, 3)
one_class_shape = (1, 21, 21, 1)
after y_true_array = (1, 21, 21, 3)
error shape =  (1, 21, 21, 3)
loss shape =  () tf.Tensor(437.87784, shape=(), dtype=float32)
TensorShape([])

When running the model, I get an InvalidArgumentError and none of the debug print statements from the loss function are printed.

model.fit(x = train_generator, steps_per_epoch = 1, epochs=2, batch_size=batch_size, shuffle=False)

Error:

InvalidArgumentError:  Input to reshape is a tensor with 4 values, but the requested shape has 1
     [[node mobilenet_loss/Reshape (defined at <ipython-input-41-4e7043bd8d70>:21) ]] [Op:__inference_train_function_12994]

Function call stack:
train_function

Please help me with this error. I'm a newbie and want to implement more models to gain experience