def METHOD_NAME(self, sample: dict[str, Any]) -> dict[str, Any]:
    """Apply the augmentation.
    Args:
        sample: Input sample.
    Returns:
        Augmented sample.
    """
    for key in ["image", "mask"]:
        dtype = sample[key].dtype
        # All inputs must be float
        sample[key] = sample[key].float()
        sample[key] = self.aug(sample[key])
        sample[key] = sample[key].to(dtype)
        # Kornia adds batch dimension
        sample[key] = rearrange(sample[key], "() c h w -> c h w")
    return sample