def _update_augmentations(self) -> None:
        """Update the augmentations for each subset."""
        for subset_name in ["train", "val", "test"]:
            subset = getattr(self, f"{subset_name}_data", None)
            augmentations = getattr(self, f"{subset_name}_augmentations", None)
            model_transform = get_nested_attr(self, "trainer.model.pre_processor.transform")
            if subset and model_transform:
                self._update_subset_augmentations(subset, augmentations, model_transform)
Update the augmentations for each subset.
src/anomalib/data/datamodules/base/image.py