RandomMosaic not working with masks?

[Sapf3ar]

Describe the bug

/.conda/lib/python3.10/site-packages/kornia/augmentation/_2d/mix/base.py", line 124, in apply_non_transform_mask raise NotImplementedError NotImplementedError

Reproduction steps

1. mosaic_mixup = kornia.augmentation.RandomMosaic(data_keys=['input','mask','mask']
2. input_shape = [4, 640, 640], mask_shape = [4,640, 640], [4, 640, 640]

Expected behavior

Cut out masks to match image crops and compose into mosaic augmentation

Environment

PyTorch version: 2.0.1+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.29.3
Libc version: glibc-2.31

Additional context

No response

[edgarriba]

@shijianjian @johnnv1 any ideas here ?

[johnnv1]

We don't have it implemented mix aug for masks, each augmentation should need a specific implementation and the RandomMosaic doesn't have it... open to contributions 😄