From 53ec6b833ed9ccb611d53510e23a58bed45b8dcd Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Tue, 22 Sep 2020 18:45:05 +0800
Subject: [PATCH] fix(mge): fix graph binding refcnt

GitOrigin-RevId: c955eb260ac3d366cbdb8761e1fef3b12449f6f6
---
 imperative/python/src/graph_rt.h            | 2 +-
 imperative/python/test/unit/test_cgtools.py | 9 +++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/imperative/python/src/graph_rt.h b/imperative/python/src/graph_rt.h
index ee2a11da8..c069e3a41 100644
--- a/imperative/python/src/graph_rt.h
+++ b/imperative/python/src/graph_rt.h
@@ -25,7 +25,7 @@ class GraphNodePtr {
     T* m_node;
 public:
     GraphNodePtr(T* node) :
-        m_graph(node ? nullptr : node->owner_graph()->shared_from_this()),
+        m_graph(node ? node->owner_graph()->shared_from_this() : nullptr),
         m_node(node) {}
     T* operator->() {return m_node;}
     T& operator*() {return *m_node;}
diff --git a/imperative/python/test/unit/test_cgtools.py b/imperative/python/test/unit/test_cgtools.py
index e74f2db17..3f0f341e1 100644
--- a/imperative/python/test/unit/test_cgtools.py
+++ b/imperative/python/test/unit/test_cgtools.py
@@ -89,3 +89,12 @@ def test_graph_traversal():
     _, var_idx = var2oprs[input_var.id][0]
 
     assert var_idx == 0
+
+
+def test_load_refcnt():
+    graph = mgb_graph.Graph()
+    varnode = graph.make_const(0)
+    buf, _ = mgb_graph.dump_graph([varnode])
+    graph, _, (varnode,) = mgb_graph.load_graph(io.BytesIO(buf))
+    del graph
+    varnode.owner
-- 
GitLab