compatible with new config

open-mmlab · Jan 22, 2024 · a3ffe2e · a3ffe2e
1 parent 17a886c
commit a3ffe2e
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/mmpretrain/apis/model.py b/mmpretrain/apis/model.py
@@ -1,5 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import copy
+import inspect
 import fnmatch
 import os.path as osp
 import re
@@ -207,8 +208,12 @@ def get_model(model: Union[str, Config],
             dataset_meta = {'classes': checkpoint['meta']['CLASSES']}
 
     if len(dataset_meta) == 0 and 'test_dataloader' in config:
-        from mmpretrain.registry import DATASETS
-        dataset_class = DATASETS.get(config.test_dataloader.dataset.type)
+        # compatible with new config
+        if inspect.isclass(config.test_dataloader.dataset.type):
+            dataset_class = config.test_dataloader.dataset.type
+        else:
+            from mmpretrain.registry import DATASETS
+            dataset_class = DATASETS.get(config.test_dataloader.dataset.type)
         dataset_meta = getattr(dataset_class, 'METAINFO', {})
 
     if device_map is not None: