Internal change

PiperOrigin-RevId: 365914010

official/vision/beta/ops/spatial_transform_ops.py

@@ -523,7 +523,7 @@ def nearest_upsampling(data, scale):
     w = shape[2]
     bs = -1 if bs is None else bs
     # Uses reshape to quickly upsample the input.  The nearest pixel is selected
-    # implicitly via broadcasting.
-    data = tf.reshape(data, [bs, h, 1, w, 1, c]) * tf.ones(
-        [1, 1, scale, 1, scale, 1], dtype=data.dtype)
+    # via tiling.
+    data = tf.tile(
+        tf.reshape(data, [bs, h, 1, w, 1, c]), [1, 1, scale, 1, scale, 1])
     return tf.reshape(data, [bs, h * scale, w * scale, c])