diff --git a/mmseg/apis/inference.py b/mmseg/apis/inference.py
index aab11d14f4..62ba264af4 100644
--- a/mmseg/apis/inference.py
+++ b/mmseg/apis/inference.py
@@ -58,11 +58,10 @@ def init_model(config: Union[str, Path, Config],
     model = MODELS.build(config.model)
     if checkpoint is not None:
         checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
-        dataset_meta = checkpoint['meta'].get('dataset_meta', None)
         # save the dataset_meta in the model for convenience
         if 'dataset_meta' in checkpoint.get('meta', {}):
             # mmseg 1.x
-            model.dataset_meta = dataset_meta
+            model.dataset_meta = checkpoint['meta']['dataset_meta']
         elif 'CLASSES' in checkpoint.get('meta', {}):
             # < mmseg 1.x
             classes = checkpoint['meta']['CLASSES']