Repository: Project-MONAI/MONAI

Issue:
Should detach the Tensor before decollate_batch
**Describe the bug**
Currently, we decollate a Tensor data(like: pred) from batch to list without `detach()` first.
I will enhance it soon.




Discussion:

Patch:
diff --git a/monai/data/utils.py b/monai/data/utils.py
--- a/monai/data/utils.py
+++ b/monai/data/utils.py
@@ -287,7 +287,7 @@ def list_data_collate(batch: Sequence):
         raise TypeError(re_str)
 
 
-def decollate_batch(batch, detach=True):
+def decollate_batch(batch, detach: bool = True):
     """De-collate a batch of data (for example, as produced by a `DataLoader`).
 
     Returns a list of structures with the original tensor's 0-th dimension sliced into elements using `torch.unbind`.
diff --git a/monai/transforms/inverse_batch_transform.py b/monai/transforms/inverse_batch_transform.py
--- a/monai/transforms/inverse_batch_transform.py
+++ b/monai/transforms/inverse_batch_transform.py
@@ -63,6 +63,7 @@ def __init__(
         loader: TorchDataLoader,
         collate_fn: Optional[Callable] = no_collation,
         num_workers: Optional[int] = 0,
+        detach: bool = True,
     ) -> None:
         """
         Args:
@@ -74,16 +75,20 @@ def __init__(
                 default to 0 as only run 1 iteration and multi-processing may be even slower.
                 if the transforms are really slow, set num_workers for multi-processing.
                 if set to `None`, use the `num_workers` of the transform data loader.
+            detach: whether to detach the tensors. Scalars tensors will be detached into number types
+                instead of torch tensors.
+
         """
         self.transform = transform
         self.batch_size = loader.batch_size
         self.num_workers = loader.num_workers if num_workers is None else num_workers
         self.collate_fn = collate_fn
+        self.detach = detach
         self.pad_collation_used = loader.collate_fn.__doc__ == pad_list_data_collate.__doc__
 
     def __call__(self, data: Dict[str, Any]) -> Any:
 
-        decollated_data = decollate_batch(data)
+        decollated_data = decollate_batch(data, detach=self.detach)
         inv_ds = _BatchInverseDataset(decollated_data, self.transform, self.pad_collation_used)
         inv_loader = DataLoader(
             inv_ds, batch_size=self.batch_size, num_workers=self.num_workers, collate_fn=self.collate_fn
diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py
--- a/monai/transforms/post/dictionary.py
+++ b/monai/transforms/post/dictionary.py
@@ -347,15 +347,17 @@ class Decollated(MapTransform):
     Note that unlike most MapTransforms, this will decollate all data, so keys are not needed.
 
     Args:
-        batch_size: if not supplied, we try to determine it based on array lengths. Will raise an error if
-            it fails to determine it automatically.
+        detach: whether to detach the tensors. Scalars tensors will be detached into number types
+            instead of torch tensors.
+
     """
 
-    def __init__(self, keys="") -> None:
+    def __init__(self, keys="", detach: bool = True) -> None:
         super().__init__(keys=keys)
+        self.detach = detach
 
     def __call__(self, data: dict):
-        return decollate_batch(data)
+        return decollate_batch(data, detach=self.detach)
 
 
 class ProbNMSd(MapTransform):


Test patch:
diff --git a/tests/test_decollate.py b/tests/test_decollate.py
--- a/tests/test_decollate.py
+++ b/tests/test_decollate.py
@@ -126,7 +126,7 @@ def check_decollate(self, dataset):
 
         for b, batch_data in enumerate(loader):
             decollated_1 = decollate_batch(batch_data)
-            decollated_2 = Decollated()(batch_data)
+            decollated_2 = Decollated(detach=True)(batch_data)
 
             for decollated in [decollated_1, decollated_2]:
                 for i, d in enumerate(decollated):
diff --git a/tests/test_inverse.py b/tests/test_inverse.py
--- a/tests/test_inverse.py
+++ b/tests/test_inverse.py
@@ -677,7 +677,7 @@ def test_inverse_inferred_seg(self, extra_transform):
         self.assertEqual(inv_seg.shape[1:], test_data[0]["label"].shape)
 
         # Inverse of batch
-        batch_inverter = BatchInverseTransform(transforms, loader, collate_fn=no_collation)
+        batch_inverter = BatchInverseTransform(transforms, loader, collate_fn=no_collation, detach=True)
         with allow_missing_keys_mode(transforms):
             inv_batch = batch_inverter(segs_dict)
         self.assertEqual(inv_batch[0]["label"].shape[1:], test_data[0]["label"].shape)