Update dense_image_warp.py

Author: faizan-m

tensorflow_addons/image/dense_image_warp.py

@@ -64,28 +64,20 @@ def _interpolate_bilinear_with_checks(
     grid_shape = tf.shape(grid)
     query_shape = tf.shape(query_points)
 
-    with tf.control_dependencies(
-        [
-            tf.Assert(tf.equal(tf.rank(grid), 4), ["Grid must be 4D Tensor"]),
-            tf.Assert(
-                tf.greater_equal(grid_shape[1], 2), ["Grid height must be at least 2."]
-            ),
-            tf.Assert(
-                tf.greater_equal(grid_shape[2], 2), ["Grid width must be at least 2."]
-            ),
-            tf.Assert(
-                tf.equal(tf.rank(query_points), 3),
-                ["Query points must be 3 dimensional."],
-            ),
-            tf.Assert(
-                tf.equal(query_shape[2], 2), ["Query points last dimension must be 2."]
-            ),
-        ]
-    ):
+    with tf.control_dependencies([
+      tf.debugging.assert_equal(tf.rank(grid), 4, "Grid must be 4D Tensor"),
+      tf.debugging.assert_greater_equal(grid_shape[1], 2,
+                                        "Grid height must be at least 2."),
+      tf.debugging.assert_greater_equal(grid_shape[2], 2,
+                                        "Grid width must be at least 2."),
+      tf.debugging.assert_equal(
+          tf.rank(query_points), 3, "Query points must be 3 dimensional."),
+      tf.debugging.assert_equal(query_shape[2], 2,
+                                "Query points last dimension must be 2.")
+    ]):
         return _interpolate_bilinear_impl(grid, query_points, indexing, name)
 
 
-@tf.function
 def _interpolate_bilinear_impl(
     grid: types.TensorLike,
     query_points: types.TensorLike,