Update augvit_model.py

augvit_model.py

@@ -179,6 +179,8 @@ class AugViTForImageClassification(TFPreTrainedModel):
         return_dict: Optional[bool] = None,
         training: Optional[bool] = False,
         **kwargs):
-        print(pixel_values)
-        logits = self.model(pixel_values['pixel_values'])
+        inp = pixel_values['pixel_values']
+        if inp.shape[-1]!=3:
+            inp = tf.transpose(inp,[0,2,3,1])
+        logits = self.model()
         return logits