From 205b22406a502b3664e7dabdbaa59c6d2662f3c9 Mon Sep 17 00:00:00 2001
From: pintaoz <pintaoz@amazon.com>
Date: Wed, 5 Feb 2025 13:50:01 -0800
Subject: [PATCH] Add docstring for image_uris.retrieve

---
 src/sagemaker/image_uris.py | 2 ++
 1 file changed, 2 insertions(+)

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).