src/transformers/models/vit/modeling_flax_vit.py/FlaxViTForImageClassificationModule/setup
class FlaxViTForImageClassificationModule:    def setup(self):
        self.vit = FlaxViTModule(config=self.config, dtype=self.dtype, add_pooling_layer=False)
        self.classifier = nn.Dense(
            self.config.num_labels,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.variance_scaling(
                self.config.initializer_range**2, "fan_in", "truncated_normal"
            ),
        )