diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py index 859a08ceec..5bf699beea 100644 --- a/src/sagemaker/estimator.py +++ b/src/sagemaker/estimator.py @@ -85,7 +85,7 @@ def __init__(self, role, train_instance_count, train_instance_type, if self.train_instance_type == 'local_gpu' and self.train_instance_count > 1: raise RuntimeError("Distributed Training in Local GPU is not supported") - self.sagemaker_session = LocalSession() + self.sagemaker_session = sagemaker_session or LocalSession() else: self.local_mode = False self.sagemaker_session = sagemaker_session or Session()