Save and load ViT model into a unique .h5 file (or TensorflowLight)

🤗Transformers
1

Hi there, I saved the model using the following code, that load a pre-trained model, save the model, and then load/save the model in the TF version:

from transformers import AutoModelForImageClassification
model = AutoModelForImageClassification.from_pretrained('google/vit-base-patch16-224')
model.save_pretrained("my_model")

from transformers import TFAutoModelForImageClassification

tf_model = TFAutoModelForImageClassification.from_pretrained("my_model", from_pt=True)
tf_model.save_pretrained("my_model_tf")

In my_model_tf there are 2 different file:

  1. config.json, that contains the “characteristics” of the model
  2. tf_model.h5, that contains the weights of the model

When i try to load tf_model.h5, I have the following Error:

ValueError: No model config found in the file at <tensorflow.python.platform.gfile.GFile object at 0x7f83009fb110>

What i am trying to do is having ONLY ONE file .h5 (or even TenforFlowLIte) that contains my model that i can load and get the whole model!