zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-15 01:19:21 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix py gc edge case (#2079)

Browse Source
This commit is contained in:
Awni Hannun
2025-04-18 12:46:53 -07:00
committed by GitHub
parent b529515eb1
commit 55935ccae7
3 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/src/transforms.cpp
View File

@@ -960,6 +960,11 @@ class PyCustomFunction {
};
int py_custom_function_tp_traverse(PyObject* self, visitproc visit, void* arg) {
Py_VISIT(Py_TYPE(self));
if (!nb::inst_ready(self)) {
return 0;
}
auto* p = nb::inst_ptr<PyCustomFunction>(self);
nb::handle v = nb::find(p->fun_);
Py_VISIT(v.ptr());
@@ -975,7 +980,6 @@ int py_custom_function_tp_traverse(PyObject* self, visitproc visit, void* arg) {
nb::handle v = nb::find(*(p->vmap_fun_));
Py_VISIT(v.ptr());
}
Py_VISIT(Py_TYPE(self));
return 0;
}
int py_custom_function_tp_clear(PyObject* self) {