diff --git a/paddle/fluid/pybind/eager.cc b/paddle/fluid/pybind/eager.cc index c600844596d98f023a1b372b988ae01fe692e58f..74d15b6c0ca7957d2bff7a46358323686efaa556 100644 --- a/paddle/fluid/pybind/eager.cc +++ b/paddle/fluid/pybind/eager.cc @@ -709,6 +709,8 @@ int TensorInit(PyObject* self, PyObject* args, PyObject* kwargs) { } static void TensorDealloc(TensorObject* self) { + if (self->weakrefs != NULL) + PyObject_ClearWeakRefs(reinterpret_cast(self)); self->tensor.~Tensor(); Py_TYPE(self)->tp_free(reinterpret_cast(self)); } @@ -739,6 +741,7 @@ void BindEager(pybind11::module* module) { type->tp_getset = variable_properties; type->tp_init = TensorInit; type->tp_new = TensorNew; + type->tp_weaklistoffset = offsetof(TensorObject, weakrefs); Py_INCREF(&PyBaseObject_Type); type->tp_base = reinterpret_cast(&PyBaseObject_Type); type->tp_flags |= diff --git a/paddle/fluid/pybind/eager.h b/paddle/fluid/pybind/eager.h index bb55ef62ee6895a4c60b4780b039af6f1f498653..03676a677ac90255d6d97cb5427000a85bc817c0 100644 --- a/paddle/fluid/pybind/eager.h +++ b/paddle/fluid/pybind/eager.h @@ -22,6 +22,8 @@ namespace pybind { typedef struct { PyObject_HEAD paddle::experimental::Tensor tensor; + // Weak references + PyObject* weakrefs; } TensorObject; typedef struct {