From e31c60cfdd21f2a967f5eb901498c0794b097561 Mon Sep 17 00:00:00 2001 From: kaczmarj Date: Wed, 21 Jun 2023 21:17:45 -0400 Subject: [PATCH] add type annotation --- wsinfer_zoo/cli.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/wsinfer_zoo/cli.py b/wsinfer_zoo/cli.py index bd8288c..65a5eda 100644 --- a/wsinfer_zoo/cli.py +++ b/wsinfer_zoo/cli.py @@ -14,6 +14,7 @@ from wsinfer_zoo.client import HF_WEIGHTS_SAFETENSORS_NAME from wsinfer_zoo.client import InvalidModelConfiguration from wsinfer_zoo.client import InvalidRegistryConfiguration +from wsinfer_zoo.client import Model from wsinfer_zoo.client import load_registry from wsinfer_zoo.client import validate_config_json @@ -96,6 +97,7 @@ def get(*, model_name: str, weights_format: str, registry_file: str): registered_model = registry.get_model_by_name(model_name) + model: Model if weights_format == "torchscript": model = registered_model.load_model_torchscript() elif weights_format == "pytorch":