Apply the exact same transformation to two images using Torchio

I want to apply the exact same transformation to two images (image and segmentation data) using torchio. Both these images are stored in numpy arrays called image_data and segmentation_data.

So far, I added some augmentations:

self.augmentations = tio.Compose([
            affine_transform,
            elastic_transform,
            flip_transform,
            swap_transform
        ])

where e.g. elastic_transform = tio.RandomElasticDeformation and tried to apply these to the images in the following way:

        subject_image = tio.Subject(image=tio.ScalarImage(tensor=image_data))
        subject_segmentation = tio.Subject(
            image=tio.ScalarImage(tensor=segmentation_data))
        dataset = tio.SubjectsDataset([subject_image, subject_segmentation])
        dataset = self.augmentations(dataset)
        image_data = dataset[0]['image'].data
        segmentation_data = dataset[1]['image'].data

Unfortunately, that's incorrect (beacuse Compose won't work with a SubjectsDataset). How to do it correctly?

Both these tensors need to be added to the subject. Here is the correct code:

    subject = tio.Subject(image=tio.ScalarImage(tensor=image_data),
                          segmentation=tio.ScalarImage(tensor=segmentation_data))
    transformed_subject = self.augmentations(subject)
    transformed_image = transformed_subject.image.data
    transformed_segmentation = transformed_subject.segmentation.data