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

Author: pmeier

torchvision/prototype/transforms/_augment.py

@@ -6,7 +6,7 @@
 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 @@ def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
             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]))