From 59449dbafe72b5b2fb1267bb255291154e351592 Mon Sep 17 00:00:00 2001 From: bhack Date: Thu, 24 Feb 2022 00:52:47 +0100 Subject: [PATCH] add explict arg keyword --- keras_cv/layers/preprocessing/grid_mask.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/keras_cv/layers/preprocessing/grid_mask.py b/keras_cv/layers/preprocessing/grid_mask.py index 99bb9f893d..c5d69bf953 100644 --- 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)