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

Fixes segfault when compiling checkpointed functions (#1235)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-06-26 16:14:45 -07:00
committed by GitHub
parent 2615660e62
commit b05bcfd27f
3 changed files with 16 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/src/transforms.cpp
View File

@@ -536,7 +536,6 @@ struct PyCompiledFun {
class PyCheckpointedFun {
public:
PyCheckpointedFun(nb::callable fun) : fun_(std::move(fun)) {}
~PyCheckpointedFun() {
nb::gil_scoped_acquire gil;
@@ -968,5 +967,8 @@ void init_transforms(nb::module_& m) {
// Register static Python object cleanup before the interpreter exits
auto atexit = nb::module_::import_("atexit");
atexit.attr("register")(nb::cpp_function([]() { tree_cache().clear(); }));
atexit.attr("register")(nb::cpp_function([]() {
tree_cache().clear();
detail::compile_clear_cache();
}));
}