add explict arg keyword

keras-team · Feb 23, 2022 · 59449db · 59449db
1 parent 50c7d07
commit 59449db
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/keras_cv/layers/preprocessing/grid_mask.py b/keras_cv/layers/preprocessing/grid_mask.py
@@ -213,7 +213,7 @@ def _augment_images(self, images):
             images = tf.expand_dims(images, axis=0)
 
         # TODO: Make the batch operation vectorize.
-        output = tf.vectorized_map(lambda image: self._grid_mask(image), images, True)
+        output = tf.vectorized_map(lambda image: self._grid_mask(image), images, fallback_to_while_loop=True)
 
         if unbatched:
             output = tf.squeeze(output, axis=0)