Repository: kornia/kornia

Issue:
RandomAffine3D throws RuntimeError when batch_size > 1
### Describe the bug

When the batch_size of input > 1, `kornia.augmentation.RandomAffine3D` throws RuntimeError: "The size of tensor a (3) must match the size of tensor b (6) at non-singleton dimension 2".


### Reproduction steps

```bash
import torch
import kornia as K

rng = torch.manual_seed(0)
input = torch.rand(2, 1, 3, 3, 3)
aug = K.augmentation.RandomAffine3D((15., 20., 20.), p=1., return_transform=True)
aug(input)
```


### Expected behavior

We should expect the valid results.

### Environment

```shell
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```
- PyTorch Version (e.g., 1.0): 1.7.0
- OS (e.g., Linux): Linux
- How you installed PyTorch (`conda`, `pip`, source): pip
- Build command you used (if compiling from source):
- Python version: 3.7.3
- CUDA/cuDNN version: 10.1
- GPU models and configuration:
- Any other relevant information:
```


### Additional context

_No response_


Discussion:
I can reproduce it. It comes from https://github.com/kornia/kornia/blob/1c9491a850d21f5ec0f1dd8343404a79c34f942c/kornia/geometry/transform/projwarp.py#L135-L136

Probably update the code to:
```python
scaling_matrix: torch.Tensor = K.eye_like(3, scales) 
```

Can you give it a shot?

Resolving patch:
diff --git a/kornia/augmentation/random_generator/random_generator3d.py b/kornia/augmentation/random_generator/random_generator3d.py
--- a/kornia/augmentation/random_generator/random_generator3d.py
+++ b/kornia/augmentation/random_generator/random_generator3d.py
@@ -119,7 +119,7 @@ def random_affine_generator3d(
             dim=1,
         )
     else:
-        scale = torch.ones(batch_size, device=device, dtype=dtype).repeat(1, 3)
+        scale = torch.ones(batch_size, device=device, dtype=dtype).reshape(batch_size, 1).repeat(1, 3)
 
     if translate is not None:
         if translate.shape != torch.Size([3]):


Test patch:
diff --git a/test/augmentation/test_augmentation_3d.py b/test/augmentation/test_augmentation_3d.py
--- a/test/augmentation/test_augmentation_3d.py
+++ b/test/augmentation/test_augmentation_3d.py
@@ -8,6 +8,7 @@
 import kornia
 import kornia.testing as utils  # test utils
 from kornia.augmentation import (
+    RandomAffine3D,
     RandomCrop,
     RandomCrop3D,
     RandomDepthicalFlip3D,
@@ -960,3 +961,26 @@ def build_input(channels, depth, height, width, bs=1, row=None, device='cpu', dt
         batch = torch.stack([image3d] * bs)
 
         return batch.to(device, dtype)
+
+
+class TestRandomAffine3D:
+
+    def test_batch_random_affine_3d(self, device, dtype):
+
+        f = RandomAffine3D((0, 0, 0), p=1., return_transform=True)  # No rotation
+        tensor = torch.tensor([[[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]]]],
+                              device=device, dtype=dtype)  # 1 x 1 x 1 x 3 x 3
+
+        expected = torch.tensor([[[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]]]],
+                                device=device, dtype=dtype)  # 1 x 1 x 1 x 3 x 3
+
+        expected_transform = torch.tensor(
+            [[[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]],
+            device=device, dtype=dtype)  # 1 x 4 x 4
+
+        tensor = tensor.repeat(5, 3, 1, 1, 1)  # 5 x 3 x 3 x 3 x 3
+        expected = expected.repeat(5, 3, 1, 1, 1)  # 5 x 3 x 3 x 3 x 3
+        expected_transform = expected_transform.repeat(5, 1, 1)  # 5 x 4 x 4
+
+        assert (f(tensor)[0] == expected).all()
+        assert (f(tensor)[1] == expected_transform).all()