Update beta sampling code in augment.py

tensorflow · Dec 11, 2024 · a2a8547 · a2a8547
1 parent 8c5c79c
commit a2a8547
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/official/vision/ops/augment.py b/official/vision/ops/augment.py
@@ -2709,8 +2709,8 @@ def distort(self, images: tf.Tensor,
 
   @staticmethod
   def _sample_from_beta(alpha, beta, shape):
-    sample_alpha = tf.random.gamma(shape, 1., beta=alpha)
-    sample_beta = tf.random.gamma(shape, 1., beta=beta)
+    sample_alpha = tf.random.gamma(shape, alpha, beta=1.0)
+    sample_beta = tf.random.gamma(shape, beta, beta=1.0)
     return sample_alpha / (sample_alpha + sample_beta)
 
   def _cutmix(self, images: tf.Tensor,