diff --git a/torchvision/prototype/transforms/__init__.py b/torchvision/prototype/transforms/__init__.py
index dc6476ab4b5..ca89fee918a 100644
--- a/torchvision/prototype/transforms/__init__.py
+++ b/torchvision/prototype/transforms/__init__.py
@@ -3,7 +3,7 @@
 from ._transform import Transform  # usort: skip
 
 from ._augment import RandomCutmix, RandomErasing, RandomMixup
-from ._auto_augment import AugMix, AutoAugment, RandAugment, TrivialAugmentWide
+from ._auto_augment import AugMix, AutoAugment, AutoAugmentPolicy, RandAugment, TrivialAugmentWide
 from ._color import (
     ColorJitter,
     RandomAdjustSharpness,