diff --git a/mmflow/apis/inference.py b/mmflow/apis/inference.py
index 42aa44c..9b1ef55 100644
--- a/mmflow/apis/inference.py
+++ b/mmflow/apis/inference.py
@@ -88,7 +88,7 @@ def inference_model(model: torch.nn.Module, img1s: Union[str, np.ndarray],
     # there is no need to load annotation.
     for t in cfg.pipeline:
         if t.get('type') == 'LoadAnnotations':
-            cfg.test_pipeline.remove(t)
+            cfg.pipeline.remove(t)
 
     test_pipeline = Compose(cfg.pipeline)
     datas = defaultdict(list)
diff --git a/mmflow/models/decoders/base_decoder.py b/mmflow/models/decoders/base_decoder.py
index 6021cc8..744f9e7 100644
--- a/mmflow/models/decoders/base_decoder.py
+++ b/mmflow/models/decoders/base_decoder.py
@@ -99,7 +99,7 @@ def postprocess_result(
                             f[1, :, :] = f[1, :, :] / h_scale
                     data_samples[i].set_data(
                         {'pred_' + key: PixelData(**{'data': f})})
-            return data_samples
+        return data_samples
 
     def predict_by_feat(self,
                         flow_results: Tensor,