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

Let pybind handle deletion of cancellation managers

Browse Source 
PiperOrigin-RevId: 355403635
Change-Id: I062d302af282e3f03c369292fffa891e54e254b2
This commit is contained in:
Allen Lavoie 2021-02-03 08:35:47 -08:00 committed by TensorFlower Gardener
parent bcc6c99c90
commit 3ef35ee4b3
2 changed files with 2 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/python/eager/cancellation.py
View File

@ -41,6 +41,3 @@ class CancellationManager(object):
def get_cancelable_function(self, concrete_function):
# pylint: disable=protected-access
return concrete_function._experimental_with_cancellation_manager(self)
def __del__(self):
pywrap_tfe.TFE_DeleteCancellationManager(self._impl)

8
tensorflow/python/tfe_wrapper.cc
View File

@ -1354,16 +1354,12 @@ PYBIND11_MODULE(_pywrap_tfe, m) {
py::return_value_policy::reference);
// TFE_CancellationManager Logic
m.def(
"TFE_NewCancellationManager",
[]() { return new tensorflow::CancellationManager(); },
py::return_value_policy::reference);
m.def("TFE_NewCancellationManager",
[]() { return new tensorflow::CancellationManager(); });
m.def("TFE_CancellationManagerIsCancelled",
&tensorflow::CancellationManager::IsCancelled);
m.def("TFE_CancellationManagerStartCancel",
&tensorflow::CancellationManager::StartCancel);
m.def("TFE_DeleteCancellationManager",
[](tensorflow::CancellationManager* cm) { delete cm; });
m.def("TFE_ClearScalarCache", &tensorflow::TFE_ClearScalarCache);