diff --git a/src/braket/jobs/image_uris.py b/src/braket/jobs/image_uris.py index c1e1457fe..eedc5e795 100644 --- a/src/braket/jobs/image_uris.py +++ b/src/braket/jobs/image_uris.py @@ -26,11 +26,11 @@ class Framework(str, Enum): PL_PYTORCH = "PL_PYTORCH" -@cache def built_in_images(region): return {retrieve_image(framework, region) for framework in Framework} +@cache def retrieve_image(framework: Framework, region: str) -> str: """Retrieves the ECR URI for the Docker image matching the specified arguments.