Update monai/transforms/post/dictionary.py

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Signed-off-by: sewon jeon <irocks0922@gmail.com>

Author: eclipse0922

monai/transforms/post/dictionary.py

@@ -693,28 +693,27 @@ def __call__(self, data: Mapping[Hashable, Any]) -> dict[Hashable, Any]:
                 # If orig_key == key, the data at d[orig_key] may have been modified by
                 # postprocessing transforms. We need to exclude any transforms that were
                 # added after the preprocessing pipeline completed.
+                # When orig_key == key, filter out postprocessing transforms to prevent
+                # confusion during inversion (see issue #8396)
                 if orig_key == key:
-                    # Count how many invertible transforms are in the preprocessing pipeline
-                    # This gives us the expected number of transforms that should be in
-                    # applied_operations from preprocessing
                     num_preproc_transforms = 0
-                    if hasattr(self.transform, 'transforms'):
-                        for t in self.transform.flatten().transforms:
-                            if isinstance(t, InvertibleTransform):
-                                num_preproc_transforms += 1
-                    elif isinstance(self.transform, InvertibleTransform):
-                        num_preproc_transforms = 1
-
-                    # Use only the first N transforms from applied_operations,
-                    # where N is the number of preprocessing transforms
-                    # This excludes any postprocessing transforms like Lambdad
+                    try:
+                        if hasattr(self.transform, 'transforms'):
+                            for t in self.transform.flatten().transforms:
+                                if isinstance(t, InvertibleTransform):
+                                    num_preproc_transforms += 1
+                        elif isinstance(self.transform, InvertibleTransform):
+                            num_preproc_transforms = 1
+                    except AttributeError:
+                        # Fallback: use all transforms if flatten fails
+                        num_preproc_transforms = len(all_transforms)
+
                     if num_preproc_transforms > 0:
                         transform_info = all_transforms[:num_preproc_transforms]
                     else:
                         transform_info = all_transforms
                 else:
                     transform_info = all_transforms
-            else:
                 transform_info = d[InvertibleTransform.trace_key(orig_key)]
                 meta_info = d.get(orig_meta_key, {})
             if nearest_interp: