[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

kornia/augmentation/auto/operations/policy.py

@@ -74,9 +74,9 @@
                     flags = override_parameters(module.op.flags, extra_args, in_place=False)
                     mat = module.op.generate_transformation_matrix(input, param.data, flags)
                 elif module.op._transform_matrix is not None:
-                    mat = as_tensor(module.op._transform_matrix, device=input.device, dtype=input.dtype)
+                    mat = as_tensor(module.transform_matrix, device=input.device, dtype=input.dtype)
                 else:
-                    raise RuntimeError(f"{module}.op._transform_matrix is None while `recompute=False`.")
+                    raise RuntimeError(f"{module}.transform_matrix is None while `recompute=False`.")
                 res_mat = mat @ res_mat
                 input = module.op.transform_output_tensor(input, ori_shape)
                 if module.op.keepdim and ori_shape != input.shape:

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)