diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index 7bb7497abc..4722727171 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -174,6 +174,10 @@ void addObjectMethodsForTraining(py::module& m) { Environment& env = get_env(); return onnxruntime::make_unique(so, env); })) + .def(py::init([]() { + Environment& env = get_env(); + return onnxruntime::make_unique(GetDefaultCPUSessionOptions(), env); + })) .def("finalize", [](py::object) { #ifdef USE_HOROVOD training::shutdown_horovod(); diff --git a/orttraining/orttraining/python/training/training_session.py b/orttraining/orttraining/python/training/training_session.py index c4dba95ab7..7357b48f94 100644 --- a/orttraining/orttraining/python/training/training_session.py +++ b/orttraining/orttraining/python/training/training_session.py @@ -13,11 +13,9 @@ from onnxruntime.capi.session import Session, InferenceSession, IOBinding class TrainingSession(InferenceSession): def __init__(self, path_or_bytes, parameters, sess_options=None): if sess_options: - self._sess = C.TrainingSession( - sess_options, C.get_session_initializer()) + self._sess = C.TrainingSession(sess_options) else: - self._sess = C.TrainingSession( - C.get_session_initializer(), C.get_session_initializer()) + self._sess = C.TrainingSession() if isinstance(path_or_bytes, str): self._sess.load_model(path_or_bytes, parameters)