from transformers import FlaxGPT2LMHeadModel

# Load model
flax_model = FlaxGPT2LMHeadModel.from_pretrained("gpt2")

# Save Flax model
flax_model.save_pretrained("LopezLLM")

print("Flax model saved!")