class_docstr: ViT model for classification task.
class_docstr_tok: ["ViT","model","for","classification","task","."]
code: def build_flax_model(self):
      return vit.ViT(
          num_classes=self.dataset_meta_data["num_classes"],
          mlp_dim=self.config.model.mlp_dim,
          num_layers=self.config.model.num_layers,
          num_heads=self.config.model.num_heads,
          representation_size=self.config.model.representation_size,
          patches=self.config.model.patches,
          hidden_size=self.config.model.hidden_size,
          classifier=self.config.model.classifier,
          dropout_rate=self.config.model.get("dropout_rate", 0.1),
          attention_dropout_rate=self.config.model.get(
              "attention_dropout_rate", 0.1
          ),
          dtype="float32",
      )
  
file_path: scenic/scenic/projects/tasseo/vit.py
imports: from typing import Optional
  from flax.training import common_utils
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.projects.baselines import vit
  from topological_transformer.images import topvit
label_desc: Use image node_id 1 for calling the ViTClassificationModel obj's underlying member method code with example usage: obj.build_flax_model() and returns: vit