fix random erase probability

official/vision/beta/ops/augment.py

@@ -1358,8 +1358,8 @@ class RandomErasing(ImageAugment):
       tf.Tensor: The augmented version of `image`.
     """
     uniform_random = tf.random.uniform(shape=[], minval=0., maxval=1.0)
-    mirror_cond = tf.less(uniform_random, .5)
-    tf.cond(mirror_cond, lambda: self._erase(image), lambda: image)
+    mirror_cond = tf.less(uniform_random, self._probability)
+    image = tf.cond(mirror_cond, lambda: self._erase(image), lambda: image)
     return image
 
   @tf.function
@@ -1399,7 +1399,7 @@ class RandomErasing(ImageAugment):
                 dtype=tf.int32)
             
             image = _fill_rectangle(image, center_width, center_height,
-                                half_width, half_height, replace=None)
+                                    half_width, half_height, replace=None)
 
             is_trial_successfull = True
 

official/vision/beta/ops/augment_test.py

@@ -263,7 +263,7 @@ class RandomErasingTest(tf.test.TestCase, parameterized.TestCase):
     aug_image = augmenter.distort(image)
 
     self.assertEqual((224, 224, 3), aug_image.shape)
-    self.assertLess(0, tf.reduce_max(aug_image.shape))
+    self.assertNotEqual(0, tf.reduce_max(aug_image))
 
 
 class MixupAndCutmixTest(tf.test.TestCase, parameterized.TestCase):