from whisper_jax import FlaxWhisperForConditionalGeneration
import jax.numpy as jnp

checkpoint_id = "/media/user01/HDWINDOWS/whisper-medium-portuguese"
# convert PyTorch weights to Flax
model = FlaxWhisperForConditionalGeneration.from_pretrained(checkpoint_id, from_pt=True)

# Save the converted Flax model in the same directory
model.save_pretrained(checkpoint_id)