[BUG] Inconsistent batch size output

Thanks for the great package. The output of kornia.augmentation.RandomRotation has an inconsistent batch size when the input batch size is 1:

import torch
from kornia.augmentation import RandomRotation  # version 0.2.0

rotation = RandomRotation((-90.,90.))

x = torch.rand(1,3,10,10)
print(rotation(x).shape)   # torch.Size([3, 10, 10])  // expecting torch.Size([1, 3, 10, 10])

x = torch.rand(2,3,10,10)
print(rotation(x).shape)   # torch.Size([2, 3, 10, 10]) // fine

x = torch.rand(6,3,10,10)
print(rotation(x).shape)   # torch.Size([6, 3, 10, 10]) // fine

Is this desired behavior?

[ducha-aiki]

I think that keeping (B, C, H, W) even for B=1 and/or C=1 is good convention to enforce.

[andraugust]

By allowing only (B, C, H, W) input shape, the output shape can be uniquely unified to (B, C, H, W).

I agree. I think input and output shape should always be (B, C, H, W), even when B=1 and/or C=1.