how to set a custom prior for Convolution2DReparameterization?

Question

I have created a Bayesian CNN and now i want to try different distributions for my prior in Convolution2DReparameterization layer and i get this error:

/usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/layers/conv_variational.py in build(self, input_shape) 190 self.kernel_prior = None 191 else: --> 192 self.kernel_prior = self.kernel_prior_fn( 193 dtype, kernel_shape, 'kernel_prior', 194 self.trainable, self.add_variable)

TypeError: Sequential.call() takes from 2 to 4 positional arguments but 6 were given

My Code:

def dist(shape):
    dist = tfd.MultivariateNormalDiag(loc=1.2 * tf.ones(shape),
                                      scale_diag=3.0*tf.ones(shape))

    batch_ndims = tf.size(dist.batch_shape_tensor())

    return tfd.Independent(dist, reinterpreted_batch_ndims=batch_ndims)



def cnn_prior(kernel_size, dtype=None):

    num_kernel_params = kernel_size[0] * kernel_size[1] * kernel_size[2] * bias_size   
    n = num_kernel_params + bias_size
    prior_model = Sequential([tfpl.DistributionLambda(lambda t : dist(kernel_size))]) 
    return prior_model



layer = tfpl.Convolution2DReparameterization(
                input_shape= input_shape, filters= filters, kernel_size= (3,3),
                activation='swish', padding='VALID',

                kernel_prior_fn= cnn_prior((3,3,1)),
                kernel_posterior_fn= tfpl.default_mean_field_normal_fn(is_singular=False),
                kernel_divergence_fn= divergence_fn,

                bias_prior_fn= cnn_prior((3,3,1)),
                bias_posterior_fn= tfpl.default_mean_field_normal_fn(is_singular=False),
                bias_divergence_fn= divergence_fn`