Update classification in converter (#4146)

openvinotoolkit · Dec 5, 2024 · 3941190 · 3941190
1 parent d9a3cd8
commit 3941190
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/otx/tools/converter.py b/src/otx/tools/converter.py
@@ -240,7 +240,7 @@ def convert(config_path: str, task: OTXTaskType | None = None) -> dict:
         if param_dict.get("enable_tiling", None) and not task_info["model_name"].endswith("_tile"):
             task_info["model_name"] += "_tile"
         # classification task type can't be deducted from template name, try to extract from config
-        if "sub_task_type" in template_config and "CLS" in template_config["sub_task_type"]:
+        if "sub_task_type" in template_config and "_CLS" in task_info["task"]:
             task_info["task"] = template_config["sub_task_type"]
         if task is not None:
             task_info["task"] = task