Changes related to TF backend compatibility check for Vision Transformer on small datasets

[pksX01]

trafficstars

This PR caters to address changes required to convert vit_small_ds.py to Keras 3 with Tensorflow backend. This PR does not add backend_agnostic changes.