From f335b610ca66c5e7bbdc986e081c2b08c65a47c2 Mon Sep 17 00:00:00 2001
From: Joris-Kuehl-TU-Berlin <kuehl@math.tu-berlin.de>
Date: Mon, 5 Aug 2024 15:59:26 +0200
Subject: [PATCH] Fix #3558

---
 mmseg/apis/inference.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

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']