From 2b6ade24b3479a641f43aec7bded20bfbb19603e Mon Sep 17 00:00:00 2001 From: nateraw Date: Mon, 13 Sep 2021 23:31:28 -0400 Subject: [PATCH] :art: write model card to enable inference --- timm/models/hub.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/timm/models/hub.py b/timm/models/hub.py index 4cdca7f0..184e9b85 100644 --- a/timm/models/hub.py +++ b/timm/models/hub.py @@ -159,6 +159,13 @@ def push_to_hf_hub( ) with repo.commit(commit_message): + # Save model weights and config save_pretrained_for_hf(model, repo.local_dir, **config_kwargs) + # Save a model card if it doesn't exist, enabling inference. + readme_path = Path(repo.local_dir) / 'README.md' + readme_txt = f'---\ntags:\n- image-classification\n- timm\nlibrary_tag: timm\n---\n# Model card for {repo_name}' + if not readme_path.exists(): + readme_path.write_text(readme_txt) + return repo.git_remote_url()