Patch AWS_DEFAULT_REGION

cohere-ai · Sep 30, 2024 · e89f543 · e89f543
1 parent 05e3841
commit e89f543
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/cohere/manually_maintained/cohere_aws/client.py b/src/cohere/manually_maintained/cohere_aws/client.py
@@ -34,7 +34,9 @@ def __init__(
         """
         self._client = boto3.client("sagemaker-runtime", region_name=aws_region)
         self._service_client = boto3.client("sagemaker", region_name=aws_region)
-        self._sess = sage.Session(boto3.session.Session())
+        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)
         self.mode = Mode.SAGEMAKER