def METHOD_NAME(self, inputs):
    images = inputs.get("images", None)
    labels = inputs.get("labels", None)
    segmentation_masks = inputs.get("segmentation_masks", None)
    if images is None or (labels is None and segmentation_masks is None):
        raise ValueError(
            "CutMix expects inputs in a dictionary with format "
            '{"images": images, "labels": labels}. or'
            '{"images": images, "segmentation_masks": segmentation_masks}. '
            f"Got: inputs = {inputs}."
        )
    if labels is not None and not labels.dtype.is_floating:
        raise ValueError(
            f"CutMix received labels with type {labels.dtype}. "
            "Labels must be of type float."
        )