def get_checkpoint_dir(output_dir, model_arch):
    """Get framework-agnostic architecture name. Used to save all PT/TF/Flax models into the same directory."""

    arch_name = model_arch.__name__
    if arch_name.startswith("TF"):
        arch_name = arch_name[2:]
    elif arch_name.startswith("Flax"):
        arch_name = arch_name[4:]

    return os.path.join(output_dir, arch_name)
Get framework-agnostic architecture name. Used to save all PT/TF/Flax models into the same directory.
dpkd/transformers/utils/create_dummy_models.py