Update monai/transforms/inverse.py

Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com> Signed-off-by: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Project-MONAI · Dec 27, 2023 · ba82172 · ba82172
1 parent c06a89b
commit ba82172
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/monai/transforms/inverse.py b/monai/transforms/inverse.py
@@ -188,9 +188,14 @@ def track_transform_meta(
             try:
                 affine = orig_affine @ to_affine_nd(len(orig_affine) - 1, affine, dtype=torch.float64)
             except RuntimeError as e:
-                raise RuntimeError(
-                    "mismatch affine matrix, ensured that the batch dimension is not included in the calculation."
-                ) from e
+                if orig_affine.ndim > 2:
+                    if data_t.is_batch:
+                        msg = "Transform applied to batched tensor, should be applied to instances only"
+                    else:
+                        msg = "Mismatch affine matrix, ensured that the batch dimension is not included in the calculation."
+                    raise RuntimeError(msg) from e
+                else:
+                    raise
             out_obj.meta[MetaKeys.AFFINE] = convert_to_tensor(affine, device=torch.device("cpu"), dtype=torch.float64)
 
         if not (get_track_meta() and transform_info and transform_info.get(TraceKeys.TRACING)):