diff options
author | 2018-03-26 11:47:50 -0700 | |
---|---|---|
committer | 2018-03-26 11:50:29 -0700 | |
commit | d2604f8dcb8a63ca063f712c24ce5aa63403b0aa (patch) | |
tree | 6109456e2238ab20f647639aa8f05c0aba5b128d /tensorflow/contrib/lite/kernels/eigen_support.cc | |
parent | 6d46c21e9f300d07e30a2185671f07d34fac3999 (diff) |
Revert to initializing number of threads when SetNumThreads is called. Requiring it
to happen before OpInit() is way too confusing for users.
PiperOrigin-RevId: 190499644
Diffstat (limited to 'tensorflow/contrib/lite/kernels/eigen_support.cc')
-rw-r--r-- | tensorflow/contrib/lite/kernels/eigen_support.cc | 7 |
1 files changed, 7 insertions, 0 deletions
diff --git a/tensorflow/contrib/lite/kernels/eigen_support.cc b/tensorflow/contrib/lite/kernels/eigen_support.cc index 213e465552..f1fdb42624 100644 --- a/tensorflow/contrib/lite/kernels/eigen_support.cc +++ b/tensorflow/contrib/lite/kernels/eigen_support.cc @@ -46,8 +46,15 @@ void DecrementUsageCounter(TfLiteContext* context) { } if (--ptr->num_references == 0) { delete ptr; + context->eigen_context = nullptr; } } +void SetNumThreads(TfLiteContext* context, int num_threads) { + IncrementUsageCounter(context); + Eigen::setNbThreads(num_threads); + DecrementUsageCounter(context); +} + } // namespace eigen_support } // namespace tflite |