Modify `augment.transform` to use `transforms.shape.rank` instead of `tf.rank`.

PiperOrigin-RevId: 313435294

official/vision/image_classification/augment.py

@@ -169,7 +169,7 @@ def transform(image: tf.Tensor, transforms) -> tf.Tensor:
   """Prepares input data for `image_ops.transform`."""
   original_ndims = tf.rank(image)
   transforms = tf.convert_to_tensor(transforms, dtype=tf.float32)
-  if tf.rank(transforms) == 1:
+  if transforms.shape.rank == 1:
     transforms = transforms[None]
   image = to_4d(image)
   image = image_ops.transform(
@@ -989,7 +989,7 @@ class RandAugment(ImageAugment):
             # pylint:disable=g-long-lambda
             lambda selected_func=func, selected_args=args: selected_func(
                 image, *selected_args)))
-            # pylint:enable=g-long-lambda
+        # pylint:enable=g-long-lambda
 
       image = tf.switch_case(branch_index=op_to_select,
                              branch_fns=branch_fns,