Fix ref leak in dtype.to_complex()/to_real() (#125154)

By using `Py_NewRef`

Also, wrap `THPDtype_to_real`/`THPDtype_to_complex` calls with `HANDLE_TH_ERRORS`

Add regression test for the above issues, by calling to_complex for integral dtypes, that raises an exception and by preserving reference count to the same to_complex/to_real call to detect if leak is happeneing.

Replace
```cpp
auto dtype = (PyObject*)torch::getTHPDtype(current_dtype);
Py_INCREF(dtype);
return dtype;
```
with a more compact/streamlined equivalent
```cpp
return Py_NewRef(torch::getTHPDtype(current_dtype));
```

Fixes https://github.com/pytorch/pytorch/issues/124868

Pull Request resolved: https://github.com/pytorch/pytorch/pull/125154
Approved by: https://github.com/Skylion007, https://github.com/albanD
This commit is contained in:
Nikita Shulga
2024-04-29 23:59:23 +00:00
committed by PyTorch MergeBot
parent 4d717cd7c3
commit 744f341aa4
4 changed files with 36 additions and 20 deletions

View File

@ -1048,9 +1048,7 @@ PyObject* THPModule_setFlushDenormal(PyObject* _unused, PyObject* arg) {
PyObject* THPModule_getDefaultDtype(PyObject* _unused, PyObject* arg) {
HANDLE_TH_ERRORS
auto scalar_type = torch::tensors::get_default_scalar_type();
auto dtype = (PyObject*)torch::getTHPDtype(scalar_type);
Py_INCREF(dtype);
return dtype;
return Py_NewRef(torch::getTHPDtype(scalar_type));
END_HANDLE_TH_ERRORS
}