Input to reshape is a tensor with 5242880 values, but the requested shape requires a multiple of 196608

I have the following code portion for a convolutional neural network:

class PixelShuffler(Layer):
    def __init__(self, size=(2, 2), data_format=None, **kwargs):
        super(PixelShuffler, self).__init__(**kwargs)
        self.data_format = conv_utils.normalize_data_format(data_format)
        self.size = conv_utils.normalize_tuple(size, 2, 'size')

    def call(self, inputs):

        input_shape = K.int_shape(inputs)
        
        if self.data_format == 'channels_first':
            batch_size, c, h, w = input_shape
            if batch_size is None:
                batch_size = -1
            rh, rw = self.size
            oh, ow = h * rh, w * rw
            oc = c // (rh * rw)

            out = K.reshape(inputs, (batch_size, rh, rw, oc, h, w))
            out = K.permute_dimensions(out, (0, 3, 4, 1, 5, 2))
            out = K.reshape(out, (batch_size, oc, oh, ow))
            return out

        elif self.data_format == 'channels_last':
            batch_size, h, w, c = input_shape
            if batch_size is None:
                batch_size = -1
            rh, rw = self.size
            oh, ow = h * rh, w * rw
            oc = c // (rh * rw)

            out = K.reshape(inputs, (batch_size, h, w, rh, rw, oc))
            out = K.permute_dimensions(out, (0, 1, 3, 2, 4, 5))
            out = K.reshape(out, (batch_size, oh, ow, oc))
            return out
        
    def get_config(self):
        config = {'size': self.size,
                  'data_format': self.data_format}
        base_config = super(PixelShuffler, self).get_config()

        return dict(list(base_config.items()) + list(config.items()))

And here is the model creation

I=Input((256,256,3))
x1=Conv2D(filters=16,kernel_size=(3,3),strides=(1,1),padding='SAME')(I)
x2=PixelShuffler()(x1)
x3= K.reshape(x2, (-1, 256, 256, 3))
model=Model(I,x3)

while executing the model.fit, I'm having this error.

history=model.fit(dataset,epochs=20, steps_per_epoch=(15//5)) 

InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  Input to reshape is a tensor with 5242880 values, but the requested shape requires a multiple of 196608
     [[node model/tf.reshape/Reshape (defined at \AppData\Local\Temp/ipykernel_15168/3446072390.py:1) ]]
     [[custom_loss/cond/then/_0/custom_loss/cond/cond/then/_59/custom_loss/cond/cond/remove_squeezable_dimensions/cond/pivot_t/_151/_67]]
  (1) Invalid argument:  Input to reshape is a tensor with 5242880 values, but the requested shape requires a multiple of 196608
     [[node model/tf.reshape/Reshape (defined at \AppData\Local\Temp/ipykernel_15168/3446072390.py:1) ]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_1799]

Errors may have originated from an input operation.
Input Source operations connected to node model/tf.reshape/Reshape:
 model/pixel_shuffler/Reshape_1 (defined at \anaconda3\envs\ffa\lib\site-packages\keras\backend.py:3094)

Input Source operations connected to node model/tf.reshape/Reshape:
 model/pixel_shuffler/Reshape_1 (defined at \anaconda3\envs\ffa\lib\site-packages\keras\backend.py:3094)

Function call stack:
train_function -> train_function

I am a beginner in this field and trying to implement the pixel shuffler. While upsizing the shape it throws an error for example on 512*512 but downsizing the shape for example to 128*128, it works fine.