Function collate_meta_tensor_fn

monai/data/utils.py:419–435

Collate a sequence of meta tensor into a single batched metatensor. This is called by `collage_meta_tensor` and so should not be used as a collate function directly in dataloaders.

(batch, *, collate_fn_map=None)

def collate_meta_tensor_fn(batch, *, collate_fn_map=None):
    """
    Collate a sequence of meta tensor into a single batched metatensor. This is called by `collage_meta_tensor`
    and so should not be used as a collate function directly in dataloaders.
    """
    from torch.utils.data._utils.collate import collate_tensor_fn  # imported here for pylint/mypy issues

    collated = collate_tensor_fn(batch)

    meta_dicts = [i.meta or TraceKeys.NONE for i in batch]
    common_ = set.intersection(*[set(d.keys()) for d in meta_dicts if isinstance(d, dict)])
    if common_:
        meta_dicts = [{k: d[k] for k in common_} if isinstance(d, dict) else TraceKeys.NONE for d in meta_dicts]
    collated.meta = default_collate(meta_dicts)
    collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch]
    collated.is_batch = True
    return collated


def collate_meta_tensor(batch):