Add is_multi_tensor argument to default_sam_dataset to split kwargs in loader (#1092)

Author: lufre1

micro_sam/training/training.py

@@ -582,6 +582,7 @@ def default_sam_dataset(
     min_size: int = 25,
     max_sampling_attempts: Optional[int] = None,
     rois: Optional[Union[slice, Tuple[slice, ...]]] = None,
+    is_multi_tensor: bool = True,
     **kwargs,
 ) -> Dataset:
     """Create a PyTorch Dataset for training a SAM model.
@@ -608,6 +609,7 @@ def default_sam_dataset(
         min_size: Minimal object size. Smaller objects will be filtered. By default, set to '25'.
         max_sampling_attempts: Number of sampling attempts to make from a dataset.
         rois: The region of interest(s) for the data.
+        is_multi_tensor: Whether the input data to data transforms is multiple tensors or not.
         kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 
     Returns:
@@ -684,7 +686,6 @@ def default_sam_dataset(
     if custom_label_transform is None:
         label_transform = default_label_transform
     else:
-        is_multi_tensor = kwargs.pop("is_multi_tensor", True)
         label_transform = torch_em.transform.generic.Compose(
             custom_label_transform, default_label_transform, is_multi_tensor=is_multi_tensor
         )