Fix sagemaker client (#596)

* Update client.py * Update client.py
cohere-ai · Oct 30, 2024 · f31b9a5 · f31b9a5
1 parent 8a9b556
commit f31b9a5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/cohere/manually_maintained/cohere_aws/client.py b/src/cohere/manually_maintained/cohere_aws/client.py
@@ -39,7 +39,8 @@ def __init__(
         """
         if not AWS_DEPS_AVAILABLE:
             raise CohereError("AWS dependencies not available. Please install boto3 and sagemaker.")
-        self._client = boto3.client()
+        self._client = boto3.client("sagemaker-runtime", region_name=aws_region)
+        self._service_client = boto3.client("sagemaker", region_name=aws_region)
         if os.environ.get('AWS_DEFAULT_REGION') is None:
             os.environ['AWS_DEFAULT_REGION'] = aws_region
         self._sess = sage.Session(sagemaker_client=self._service_client)