resolve comment

RunningLeon · RunningLeon · commit 6ab188f0fcc4 · 2022-03-08T19:30:12.000+08:00
diff --git a/mmdeploy/codebase/mmcls/deploy/classification_model.py b/mmdeploy/codebase/mmcls/deploy/classification_model.py
@@ -156,20 +156,20 @@ def get_classes_from_config(model_cfg: Union[str, mmcv.Config]):
     data_cfg = model_cfg.data
 
     def _get_class_names(dataset_type: str):
-        dataset = data_cfg[dataset_type]
-        if not dataset:
+        dataset = data_cfg.get(dataset_type, None)
+        if (not dataset) or (dataset.type not in module_dict):
             return None
+
         module = module_dict[dataset.type]
         if module.CLASSES is not None:
             return module.CLASSES
         return module.get_classes(dataset.get('classes', None))
 
     class_names = None
     for dataset_type in ['val', 'test', 'train']:
-        if dataset_type in data_cfg:
-            class_names = _get_class_names(dataset_type)
-            if class_names is not None:
-                break
+        class_names = _get_class_names(dataset_type)
+        if class_names is not None:
+            break
 
     if class_names is None:
         logger = get_root_logger()
