experiments
/
STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

more modifications

This commit is contained in:
Koan-Sin Tan 2020-05-20 16:33:30 +08:00
parent 30e5e29d48
commit a7c3bc9ebd
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/lite/python/interpreter.py
View File

@ -213,6 +213,8 @@ class Interpreter(object):
if num_threads:
if not isinstance(num_threads, int):
raise ValueError('type of num_threads should be int')
if num_threads < 1:
raise ValueError('num_threads should >= 1')
self._interpreter.SetNumThreads(num_threads)
# Each delegate is a wrapper that owns the delegates that have been loaded

4
tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper_pybind11.cc
View File

@ -148,8 +148,8 @@ PYBIND11_MODULE(_pywrap_tensorflow_interpreter_wrapper, m) {
)pbdoc")
.def(
"SetNumThreads",
[](InterpreterWrapper& self, int i) {
return tensorflow::PyoOrThrow(self.SetNumThreads(i));
[](InterpreterWrapper& self, int num_threads) {
return tensorflow::PyoOrThrow(self.SetNumThreads(num_threads));
},
R"pbdoc(
ask the interpreter to set the number of threads to use.