saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-23 19:32:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix python error

Browse source
This commit is contained in:
Xueyun Zhu 2020-03-25 01:59:37 +00:00
parent 8f7bd51f7a
commit e9877850a4
2 changed files with 6 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
orttraining/orttraining/python/orttraining_pybind_state.cc
View file

@ -174,6 +174,10 @@ void addObjectMethodsForTraining(py::module& m) {
Environment& env = get_env();
return onnxruntime::make_unique<onnxruntime::training::TrainingSession>(so, env);
}))
.def(py::init([]() {
Environment& env = get_env();
return onnxruntime::make_unique<onnxruntime::training::TrainingSession>(GetDefaultCPUSessionOptions(), env);
}))
.def("finalize", [](py::object) {
#ifdef USE_HOROVOD
training::shutdown_horovod();

6
orttraining/orttraining/python/training/training_session.py
View file

@ -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)