code: def augment(images, labels, aug, iters):
      """Augments both image and label. Assumes input is a tensor with
      a batch dimension and values normalized to N(0,1)."""
      IMG_MEAN = np.array(
          (104.00698793, 116.66876762, 122.67891434), dtype=np.float32
      )
  
      # Transform label shape: B, C, W, H ==> B, W, H, C
      labels_are_3d = len(labels.shape) == 4
      if labels_are_3d:
          labels = labels.permute(0, 2, 3, 1)
  
      # Transform each image independently.
      aug_images, aug_labels = [], []
      for image, label in zip(images, labels):
          # Step 1: Undo normalization transformation, convert to numpy
          image = cv2.cvtColor(
              image.numpy().transpose(1, 2, 0) + IMG_MEAN,
              cv2.COLOR_BGR2RGB,
          ).astype(np.uint8)
          label = label.numpy()
          label = label.astype("int64")
  
          # Step 2: Perform transformations on numpy images
          data = aug(image=image, mask=label)
          image, label = data["image"], data["mask"]
  
          # Step 3: Convert back to PyTorch tensors
          image = paddle.to_tensor(
              (
                  cv2.cvtColor(
                      image.astype(np.float32), cv2.COLOR_RGB2BGR
                  )
                  - IMG_MEAN
              ).transpose(2, 0, 1)
          )
          # label = np.where(label==-1, 255, label)
          label = paddle.to_tensor(label)
          if not labels_are_3d:
              label = label.astype("int64")
  
          # Add to list
          aug_images.append(image)
          aug_labels.append(label)
  
      # Stack
      images = paddle.stack(aug_images, axis=0)
      labels = paddle.stack(aug_labels, axis=0)
  
      # Transform label shape back: B, W, H, C ==> B, C, W, H
      if labels_are_3d:
          labels = labels.permute(0, 3, 1, 2)
  
      return images, labels
  
imports: import cv2
  import paddle
  import numpy
  import albumentations
label_desc: Use image node_id 2 for calling a global function with example usage: augment(images, labels, aug, iters) and returns: images,  labels