The right way to discard the classifier layer

[sayakpaul]

Hi.

I am referring to the ViT models, in particular, the AugReg notebook. I wanted to export a given ViT model to TensorFlow SavedModel without the final classification layer using the jax2tf tool. I am aware of the conversion workflow.

However, I am not sure how could I discard the final classification layer from the JAX model to export a feature extractor.

Any help is appreciated.