diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py
index 8c3449e8c4..7d277cd854 100644
--- a/src/sagemaker/image_uris.py
+++ b/src/sagemaker/image_uris.py
@@ -101,6 +101,8 @@ def retrieve(
             https://github.com/aws/deep-learning-containers/blob/master/available_images.md
             (default: None).
         distribution (dict): A dictionary with information on how to run distributed training
+        base_framework_version (str): The base version number of PyTorch or Tensorflow.
+            (default: None).
         training_compiler_config (:class:`~sagemaker.training_compiler.TrainingCompilerConfig`):
             A configuration class for the SageMaker Training Compiler
             (default: None).