diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp
index c47942b72..fcc283bba 100644
--- a/python/src/transforms.cpp
+++ b/python/src/transforms.cpp
@@ -1503,6 +1503,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([]() {
+    nb::gil_scoped_acquire gil;
+
     tree_cache().clear();
     mx::detail::compile_clear_cache();
   }));