How to use albumentations augmentation?

1

I have a augmentation function using albumentations lib.

import albumentations as albu
from mxnet.gluon.data.vision.transforms import Normalize
import albumentations as albu

def aug_image_and_keypoint(image, keypoints):
    aug_all = albu.Compose([albu.OneOf([
                    albu.Blur(blur_limit=3, always_apply=False, p=0.2),
                    albu.HorizontalFlip(always_apply=False, p=0.5),
                    albu.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.1, rotate_limit=30, interpolation=1,
                                          border_mode=4, value=None, mask_value=None, always_apply=False, p=0.2)
                    ], p=1)])
    aug = albu.Compose([aug_all], p=1, keypoint_params=albu.KeypointParams(format='xy'))
    # kps
    before_aug_keypoint = keypoints.copy()
    annotations = aug(image=image, keypoints=before_aug_keypoint)
    out_image, kps_aug = annotations["image"], annotations["keypoints"]    
    return out_image, kps_aug

I have a train_iter as

 train_iter = mx.io.ImageRecordIter(
        path_imgrec=train_data_file, 
        data_shape=(3, image_size, image_size), 
        batch_size=batch_size,
        label_width=205,
        shuffle = True,
        shuffle_chunk_size = 1024,
        seed = 1234, 
        prefetch_buffer = 10, 
        preprocess_threads = 16
    )

During training, I load the data and label as

for epoch in range(0, epoches):
        # reset data iterator
        train_iter.reset()
        batch_idx = 0
        for batch in train_iter:
            batch_idx += 1
            batch_size = batch.data[0].shape[0]
            data   = batch.data[0]
            labels = batch.label[0]
            keypoints = labels[:, 0:68*2] * image_size

How can I apply augmentation for data and keypoints using aug_image_and_keypoint on batch? Thanks