Fix error handling TestSetDefaultMobileCPUAllocator

Pull Request resolved: https://github.com/pytorch/pytorch/pull/73207
This commit is contained in:
Alban Desmaison 2022-02-22 09:52:06 -05:00 committed by PyTorch MergeBot
parent cfb6c942fe
commit 7807a83f6e

View File

@ -589,11 +589,10 @@ PyObject *THPModule_supportedQEngines(PyObject *_unused, PyObject *noargs)
{ {
auto qengines = at::globalContext().supportedQEngines(); auto qengines = at::globalContext().supportedQEngines();
auto list = THPObjectPtr(PyList_New(qengines.size())); auto list = THPObjectPtr(PyList_New(qengines.size()));
if (!list) return nullptr;
for (const auto i : c10::irange(qengines.size())) { for (const auto i : c10::irange(qengines.size())) {
PyObject *i64 = THPUtils_packInt64(static_cast<int>(qengines[i])); PyObject *i64 = THPUtils_packInt64(static_cast<int>(qengines[i]));
if (!i64) { if (!i64) return nullptr;
throw python_error();
}
PyList_SET_ITEM(list.get(), i, i64); PyList_SET_ITEM(list.get(), i, i64);
} }
return list.release(); return list.release();
@ -607,22 +606,18 @@ PyObject *THPModule_isEnabledXNNPACK(PyObject *_unused, PyObject *noargs)
PyObject *THPModule_setDefaultMobileCPUAllocator(PyObject *_unused, PyObject *noargs) PyObject *THPModule_setDefaultMobileCPUAllocator(PyObject *_unused, PyObject *noargs)
{ {
try { HANDLE_TH_ERRORS
at::globalContext().setDefaultMobileCPUAllocator(); at::globalContext().setDefaultMobileCPUAllocator();
} catch (c10::Error& e) {
THPUtils_setError(e.what());
}
Py_RETURN_NONE; Py_RETURN_NONE;
END_HANDLE_TH_ERRORS
} }
PyObject *THPModule_unsetDefaultMobileCPUAllocator(PyObject *_unused, PyObject *noargs) PyObject *THPModule_unsetDefaultMobileCPUAllocator(PyObject *_unused, PyObject *noargs)
{ {
try { HANDLE_TH_ERRORS
at::globalContext().unsetDefaultMobileCPUAllocator(); at::globalContext().unsetDefaultMobileCPUAllocator();
} catch (c10::Error& e) {
THPUtils_setError(e.what());
}
Py_RETURN_NONE; Py_RETURN_NONE;
END_HANDLE_TH_ERRORS
} }
static PyObject * THPModule_vmapmode_increment_nesting(PyObject* _unused, PyObject *arg) { static PyObject * THPModule_vmapmode_increment_nesting(PyObject* _unused, PyObject *arg) {