@patrickvonplaten I think the main purpose of this project is to do pixel-level classification or segmentation with transformer. If we can do segmentation with  ViT with then I think we should use ViT too. As far as I know, ViT is for classification not segmentation. If there is no implementation of TransUNet in JAX perhaps we can work on it? It would be amazing if we can implement it in this project…