From e0fffa5286da293601b8220f330872931d075cf5 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Tue, 12 Jan 2021 18:37:41 +0800
Subject: [PATCH] fix(mge/tensor): fix mem leak when getting value in trace

GitOrigin-RevId: 1dabd88e09fc0195eaa6cba822e6dcb4f7107dd8
---
 imperative/python/megengine/jit/tracing.py | 3 +--
 imperative/python/src/tensor.cpp           | 4 +++-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/imperative/python/megengine/jit/tracing.py b/imperative/python/megengine/jit/tracing.py
index a9c66c998..c63e3c01a 100644
--- a/imperative/python/megengine/jit/tracing.py
+++ b/imperative/python/megengine/jit/tracing.py
@@ -1002,8 +1002,7 @@ class CompiledTensorProxy:
             else:
                 # c++ will throw TraceReadError
                 return None
-            if self._isscalar:
-                self.__value = self.__value.squeeze()
+        # c++ side will handle scalar case
         return self.__value
 
     def _dev_tensor(self):
diff --git a/imperative/python/src/tensor.cpp b/imperative/python/src/tensor.cpp
index 947cb2714..bce96e15f 100644
--- a/imperative/python/src/tensor.cpp
+++ b/imperative/python/src/tensor.cpp
@@ -393,7 +393,9 @@ PyObject* TensorWrapper::numpy() {
             throw TraceReadError("value of this tensor is not read in trace");
         }
         if (m_tensor->m_flags & Tensor::Flags::SCALAR) {
-            np_val = PyArray_Squeeze(reinterpret_cast<PyArrayObject*>(np_val));
+            PyObject *np_scalar = PyArray_Squeeze(reinterpret_cast<PyArrayObject*>(np_val));
+            Py_DECREF(np_val);
+            return np_scalar;
         }
         return np_val;
     }
-- 
GitLab