Custom Loss function error in low light image enhancement model

Question

I'm trying to create a custom loss function for an image enhancement model using SSIM loss.

    def mirnet_model(num_rrg, num_mrb, channels):
    input_tensor = keras.Input(shape=[None, None, 3])
    x1 = layers.Conv2D(channels, kernel_size=(3, 3), padding="same")(input_tensor)
    for _ in range(num_rrg):
        x1 = recursive_residual_group(x1, num_mrb, channels)
    conv = layers.Conv2D(3, kernel_size=(3, 3), padding="same")(x1)
    output_tensor = layers.Add()([input_tensor, conv])
    return keras.Model(input_tensor, output_tensor)


    model = mirnet_model(num_rrg=3, num_mrb=2, channels=64)

Loss function -

   def ssim_loss(y_true, y_pred):
      return tf.reduce_mean(tf.image.ssim(y_true, y_pred, 2.0))

model training -

model.compile(
    # optimizer=optimizer, loss=charbonnier_loss, metrics=[peak_signal_noise_ratio]
    optimizer=optimizer, loss=ssim_loss, metrics=[peak_signal_noise_ratio]
)

history = model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=50,
    callbacks=[
        keras.callbacks.ReduceLROnPlateau(
            monitor="val_peak_signal_noise_ratio",
            factor=0.5,
            patience=5,
            verbose=1,
            min_delta=1e-7,
            mode="max",
            # callbacks=[cp_callback]
        )
    ],
)

I'm getting good results with charbonnier loss but it might overfit because the LOL dataset is small.

I'm getting this error with the ssim loss -

ValueError: in user code:

    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1021, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1010, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1000, in run_step  **
        outputs = model.train_step(data)
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 859, in train_step
        y_pred = self(x, training=True)
    File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/usr/local/lib/python3.7/dist-packages/keras/engine/input_spec.py", line 249, in assert_input_compatibility
        f'Input {input_index} of layer "{layer_name}" is '

    ValueError: Exception encountered when calling layer "sequential" (type Sequential).
    
    Input 0 of layer "conv2d_662" is incompatible with the layer: expected axis -1 of input shape to have value 1, but received input with shape (4, None, None, 3)
    
    Call arguments received:
      • inputs=tf.Tensor(shape=(4, None, None, 3), dtype=float32)
      • training=True
      • mask=None

How to remove this error? And also tell me if there are any good loss functions that I can use in my case.