fix bugs

keras_hub/src/models/image_classifier.py

@@ -117,10 +117,12 @@ def __init__(
                 dtype=head_dtype,
                 name="pooler",
             )
+        elif pooling == "token":
+            self.pooler = None
         else:
             raise ValueError(
                 "Unknown `pooling` type. Polling should be either `'avg'` or "
-                f"`'max'`. Received: pooling={pooling}."
+                f"`'max' or 'token'`. Received: pooling={pooling}."
             )
         self.output_dropout = keras.layers.Dropout(
             dropout,
@@ -137,7 +139,10 @@ def __init__(
         # === Functional Model ===
         inputs = self.backbone.input
         x = self.backbone(inputs)
-        x = self.pooler(x)
+        if pooling == "token":  # used for Vision Transformer(ViT)
+            x = x[:, 0]
+        else:
+            x = self.pooler(x)
         x = self.output_dropout(x)
         outputs = self.output_dense(x)
         super().__init__(

keras_hub/src/models/vit/vit_backbone.py

@@ -53,7 +53,7 @@ def __init__(
             dtype=dtype,
         )(inputs)
 
-        x = ViTEncoder(
+        output = ViTEncoder(
             num_layers=num_layers,
             num_heads=num_heads,
             hidden_dim=hidden_dim,
@@ -64,8 +64,6 @@ def __init__(
             dtype=dtype,
         )(x)
 
-        output = x[:, 0]
-
         super().__init__(
             inputs=inputs,
             outputs=output,

keras_hub/src/models/vit/vit_image_classifier.py

@@ -1,5 +1,3 @@
-import keras
-
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.image_classifier import ImageClassifier
 from keras_hub.src.models.vit.vit_backbone import ViTBackbone
@@ -12,50 +10,3 @@
 class ViTImageClassifier(ImageClassifier):
     backbone_cls = ViTBackbone
     preprocessor_cls = ViTImageClassifierPreprocessor
-
-    def __init__(
-        self,
-        backbone,
-        num_classes,
-        preprocessor=None,
-        activation=None,
-        head_dtype=None,
-        **kwargs,
-    ):
-        head_dtype = head_dtype or backbone.dtype_policy
-
-        # === Layers ===
-        self.backbone = backbone
-        self.preprocessor = preprocessor
-
-        self.output_dense = keras.layers.Dense(
-            num_classes,
-            activation=activation,
-            dtype=head_dtype,
-            name="predictions",
-        )
-
-        # === Functional Model ===
-        inputs = self.backbone.input
-        x = self.backbone(inputs)
-        outputs = self.output_dense(x)
-        super().__init__(
-            inputs=inputs,
-            outputs=outputs,
-            **kwargs,
-        )
-
-        # === Config ===
-        self.num_classes = num_classes
-        self.activation = activation
-
-    def get_config(self):
-        # Backbone serialized in `super`
-        config = super().get_config()
-        config.update(
-            {
-                "num_classes": self.num_classes,
-                "pooling": self.pooling,
-            }
-        )
-        return config