From 8cf49c42d31e6a0f93d830cc263886e67155d692 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Thu, 14 Nov 2024 11:27:19 +0100 Subject: [PATCH] More elegant solution --- src/metatrain/cli/export.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/metatrain/cli/export.py b/src/metatrain/cli/export.py index 3b8ffdcc..a27b617e 100644 --- a/src/metatrain/cli/export.py +++ b/src/metatrain/cli/export.py @@ -72,8 +72,11 @@ def _prepare_export_model_args(args: argparse.Namespace) -> None: architecture_name=architecture_name, **args.__dict__, ) - if "huggingface_api_token" in args.__dict__: - args.__dict__.pop("huggingface_api_token") + keys_to_keep = ["model", "output"] # only these are needed for `export_model`` + original_keys = list(args.__dict__.keys()) + for key in original_keys: + if key not in keys_to_keep: + args.__dict__.pop(key) def export_model(model: Any, output: Union[Path, str] = "exported-model.pt") -> None: