One can use [ViTMAEForPreTraining] for this purpose.