Update augvit_model.py

augvit_model.py

@@ -156,7 +156,7 @@ class AUGViT(Model):
 
 from transformers import TFPreTrainedModel
 from .augvit_config import AugViTConfig
-
+from typing import Dict, Optional, Tuple, Union
 class AugViTForImageClassification(TFPreTrainedModel):
     config_class = AugViTConfig
     def __init__(self, config):
@@ -173,6 +173,11 @@ class AugViTForImageClassification(TFPreTrainedModel):
             emb_dropout =config.emb_dropout
         )
 
-    def call(self, input,**kwargs):
-        logits = self.model(input)
+    def call(self, pixel_values: tf.Tensor | None = None,
+        output_hidden_states: Optional[bool] = None,
+        labels: tf.Tensor | None = None,
+        return_dict: Optional[bool] = None,
+        training: Optional[bool] = False,
+        **kwargs):
+        logits = self.model(pixel_values)
         return logits