Fix tests.

Author: datumbox

test/test_transforms_tensor.py

@@ -739,7 +739,12 @@ def test_augmix(device, fill):
     tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=device)
     batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=device)
 
-    transform = T.AugMix(fill=fill)
+    class DeterministicAugMix(T.AugMix):
+        def _sample_dirichlet(self, params: torch.Tensor) -> torch.Tensor:
+            # path the method to ensure that the order of rand calls doesn't affect the outcome
+            return params.softmax(dim=-1)
+
+    transform = DeterministicAugMix(fill=fill)
     s_transform = torch.jit.script(transform)
     for _ in range(25):
         _test_transform_vs_scripted(transform, s_transform, tensor)

torchvision/transforms/autoaugment.py

@@ -535,6 +535,10 @@ def _pil_to_tensor(self, img) -> Tensor:
     def _tensor_to_pil(self, img: Tensor):
         return F.to_pil_image(img)
 
+    def _sample_dirichlet(self, params: Tensor) -> Tensor:
+        # Must be on a separate method so that we can overwrite it in tests.
+        return torch._sample_dirichlet(params)
+
     def forward(self, orig_img: Tensor) -> Tensor:
         """
             img (PIL Image or Tensor): Image to be transformed.
@@ -560,12 +564,12 @@ def forward(self, orig_img: Tensor) -> Tensor:
 
         # Sample the beta weights for combining the original and augmented image. To get Beta, we use a Dirichlet
         # with 2 parameters. The 1st column stores the weights of the original and the 2nd the ones of augmented image.
-        m = torch._sample_dirichlet(
+        m = self._sample_dirichlet(
             torch.tensor([self.alpha, self.alpha], device=batch.device).expand(batch_dims[0], -1)
         )
 
         # Sample the mixing weights and combine them with the ones sampled from Beta for the augmented images.
-        combined_weights = torch._sample_dirichlet(
+        combined_weights = self._sample_dirichlet(
             torch.tensor([self.alpha] * self.mixture_width, device=batch.device).expand(batch_dims[0], -1)
         ) * m[:, 1].view([batch_dims[0], -1])