Using model for image processing

[dejanberic]

I have converted your model to tensorflow, but I want to use it for image denoising.
Do you have an idea of how could I use/change this model so it is viable for image denosing?

One thing that I did is to replace this line:
x = tf.math.reduce_mean(x, axis=-1)
with this line:
x = rearrange(x, 'b (x z) (y dz) -> b x y (z dz)', x = axis[0], y = axis[1], z = self.num_classes)

at the end of Perceiver model.