Update return type

openvinotoolkit · Jan 8, 2024 · b095a26 · b095a26
1 parent 486a2e3
commit b095a26
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/otx/core/model/entity/classification.py b/src/otx/core/model/entity/classification.py
@@ -29,7 +29,7 @@ class OTXMulticlassClsModel(
     """Base class for the classification models used in OTX."""
 
 
-def _create_mmpretrain_model(config: DictConfig, load_from: str) -> nn.Module:
+def _create_mmpretrain_model(config: DictConfig, load_from: str) -> tuple[nn.Module, list[str]]:
     from mmpretrain.models.utils import ClsDataPreprocessor as _ClsDataPreprocessor
     from mmpretrain.registry import MODELS
 

diff --git a/src/otx/core/utils/build.py b/src/otx/core/utils/build.py
@@ -40,7 +40,7 @@ def build_mm_model(config: DictConfig, model_registry: Registry, load_from: str
     return model
 
 
-def get_classification_layers(config: DictConfig, model_registry: Registry, prefix: str = "") -> nn.Module:
+def get_classification_layers(config: DictConfig, model_registry: Registry, prefix: str = "") -> list[str]:
     """Return classification layer names by comparing two different number of classes models."""
     sample_config = deepcopy(config)
     modify_num_classes(sample_config, 5)