diff --git a/oneflow/compatible_single_client_python/framework/session_util.py b/oneflow/compatible_single_client_python/framework/session_util.py index c0aa494a27b14e805dbccc4c6275be547c0c993d..7d928ba66cb86f24418e3999816c4db2b4500d11 100644 --- a/oneflow/compatible_single_client_python/framework/session_util.py +++ b/oneflow/compatible_single_client_python/framework/session_util.py @@ -235,6 +235,11 @@ class Session(object): if self.status_ is SessionStatus.RUNNING: self.Close() + if self.status_ != SessionStatus.CLOSED: + oneflow._oneflow_internal.ClearSessionById(self.id) + + self.status_ = SessionStatus.CLOSED + def Close(self): assert self.status_ is SessionStatus.RUNNING self.Sync() @@ -442,7 +447,7 @@ class Session(object): self.cond_var_.release() def __del__(self): - oneflow._oneflow_internal.ClearSessionById(self.id) + self.TryClose() @oneflow_export("find_or_create_module") diff --git a/oneflow/python/framework/session_util.py b/oneflow/python/framework/session_util.py index 55a4b96e7831b00bcdd41e09c1c24afee40c214b..a383a1bdd74d557556e851fac44cec33ab199bc5 100644 --- a/oneflow/python/framework/session_util.py +++ b/oneflow/python/framework/session_util.py @@ -220,6 +220,11 @@ class Session(object): if self.status_ is SessionStatus.RUNNING: self.Close() + if self.status_ != SessionStatus.CLOSED: + oneflow._oneflow_internal.ClearSessionById(self.id) + + self.status_ = SessionStatus.CLOSED + def Close(self): assert self.status_ is SessionStatus.RUNNING self.Sync() @@ -427,7 +432,7 @@ class Session(object): self.cond_var_.release() def __del__(self): - oneflow._oneflow_internal.ClearSessionById(self.id) + self.TryClose() @oneflow_export("find_or_create_module")