ValueError when applying flip augmentations to boxes

Describe the bug

Came across a strange bug while applying flip augmentations with bboxes. Running the below code snippet produces the following error some of the time.

For reference my image is of shape (B, C, H, W) and my boxes are a list of length B where each element is of shape (N,4)

ValueError: Input batch size must be the same for both tensors or 1.Got torch.Size([2, 3, 3]) and torch.Size([1, 8, 2])

The randomness of the message could be stemming from the randomness in applying the augmentation but setting p=1.0 or p=0.0 leads to no error being reported.

Reproduction steps

import kornia.augmentation as K
import torch

augs = K.AugmentationSequential(
    K.RandomHorizontalFlip(p=0.5),
    K.RandomVerticalFlip(p=0.5),
    data_keys=["input", "bbox_xyxy"],
)

image = torch.randn((2,3,200,200))
boxes = [
    torch.tensor([
        [1, 2, 2, 4],
        [2, 4, 3, 6],
    ]),
    torch.tensor([
        [1, 2, 2, 4],
    ]),
]

tfm_img, tfm_bbox = augs(image, boxes)

Expected behavior

Script should run successfully all the time

Environment

PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.5
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.23.2
Libc version: glibc-2.31

Python version: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:56:21)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-117-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.3.58
GPU models and configuration: GPU 0: NVIDIA TITAN Xp
Nvidia driver version: 470.129.06
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.6.3
[pip3] mypy==0.961
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.4
[pip3] pytorch-lightning==1.6.4
[pip3] pytorch-sphinx-theme==0.0.24
[pip3] segmentation-models-pytorch==0.2.1
[pip3] torch==1.11.0
[pip3] torchmetrics==0.9.0
[pip3] torchvision==0.12.0
[conda] blas                      2.115                       mkl    conda-forge
[conda] blas-devel                3.9.0            15_linux64_mkl    conda-forge
[conda] cudatoolkit               11.5.1              h59c8dcf_10    conda-forge
[conda] efficientnet-pytorch      0.6.3                    pypi_0    pypi
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            15_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            15_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            15_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            15_linux64_mkl    conda-forge
[conda] mkl                       2022.1.0           h84fe81f_915    conda-forge
[conda] mkl-devel                 2022.1.0           ha770c72_916    conda-forge
[conda] mkl-include               2022.1.0           h84fe81f_915    conda-forge
[conda] numpy                     1.22.4           py39hc58783e_0    conda-forge
[conda] pytorch                   1.11.0          py3.9_cuda11.5_cudnn8.3.2_0    pytorch
[conda] pytorch-lightning         1.6.4                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch-sphinx-theme      0.0.24                   pypi_0    pypi
[conda] segmentation-models-pytorch 0.2.1                    pypi_0    pypi
[conda] torchmetrics              0.9.0                    pypi_0    pypi
[conda] torchvision               0.12.0               py39_cu115    pytorch

[ashnair1]

Found the issue and opened a PR to fix this. Problem was caused by not indexing the transform matrix while applying transforms.

[ashnair1]

Technically this is still a problem and applying distinct augmentations on batch elements should be possible.