src/transformers/models/vit/modeling_flax_vit.py/FlaxViTPreTrainedModel/__init__
class FlaxViTPreTrainedModel:    def __init__(self, config: ViTConfig, input_shape=None, seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs):
        module = self.module_class(config=config, dtype=dtype, **kwargs)
        if input_shape is None:
            input_shape = (1, config.image_size, config.image_size, 3)
        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)