[Fix] Transformation Matrix Mis-calculation for autoaugmentations

kornia/augmentation/auto/base.py

@@ -92,8 +92,25 @@ def forward_parameters(self, batch_shape: torch.Size) -> List[ParamItem]:
         return params
 
     def transform_inputs(self, input: Tensor, params: List[ParamItem], extra_args: Dict[str, Any] = {}) -> Tensor:
+        for param in params:
+            module = self.get_submodule(param.name)
+            input = InputSequentialOps.transform(input, module=module, param=param, extra_args=extra_args)
+        return input
+
+    def forward(
+        self, input: Tensor, params: Optional[List[ParamItem]] = None, extra_args: Dict[str, Any] = {}
+    ) -> Tensor:
+        self.clear_state()
+
+        if params is None:
+            inp = input
+            _, out_shape = self.autofill_dim(inp, dim_range=(2, 4))
+            params = self.forward_parameters(out_shape)
+
         for param in params:
             module = self.get_submodule(param.name)
             input = InputSequentialOps.transform(input, module=module, param=param, extra_args=extra_args)
             self._update_transform_matrix_by_module(module)
+
+        self._params = params
         return input

tests/augmentation/test_auto_operation.py

@@ -11,7 +11,6 @@
 from kornia.augmentation.auto.trivial_augment import TrivialAugment
 from kornia.augmentation.container import AugmentationSequential
 from kornia.geometry.bbox import bbox_to_mask
-from kornia.utils._compat import torch_version
 
 from testing.base import BaseTester
 
@@ -112,10 +111,6 @@ def test_smoke(self, policy):
         in_tensor = torch.rand(10, 3, 50, 50, requires_grad=True)
         aug(in_tensor)
 
-    @pytest.mark.xfail(
-        torch_version() in {"1.9.1", "1.10.2", "1.11.0", "1.12.1", "1.13.1"},
-        reason="randomness failing into some torch versions",
-    )
     def test_transform_mat(self, device, dtype):
         aug = RandAugment(n=3, m=15)
         in_tensor = torch.rand(10, 3, 50, 50, device=device, dtype=dtype, requires_grad=True)