add support for apply probability to CutMix and MixUp (#6448)

torchvision/prototype/transforms/_augment.py

@@ -6,7 +6,7 @@ from typing import Any, Dict, Tuple
 import PIL.Image
 import torch
 from torchvision.prototype import features
-from torchvision.prototype.transforms import functional as F, Transform
+from torchvision.prototype.transforms import functional as F
 
 from ._transform import _RandomApplyTransform
 from ._utils import get_image_dimensions, has_all, has_any, is_simple_tensor, query_image
@@ -97,9 +97,9 @@ class RandomErasing(_RandomApplyTransform):
             return inpt
 
 
-class _BaseMixupCutmix(Transform):
-    def __init__(self, *, alpha: float) -> None:
-        super().__init__()
+class _BaseMixupCutmix(_RandomApplyTransform):
+    def __init__(self, *, alpha: float, p: float = 0.5) -> None:
+        super().__init__(p=p)
         self.alpha = alpha
         self._dist = torch.distributions.Beta(torch.tensor([alpha]), torch.tensor([alpha]))