[Bugfix] Fix performance bug in edge softmax operator (dmlc#624)

gwli · Jun 7, 2019 · 8bf9771 · 8bf9771
1 parent 7ad0485
commit 8bf9771
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/dgl/nn/pytorch/softmax.py b/python/dgl/nn/pytorch/softmax.py
@@ -126,7 +126,8 @@ def backward(ctx, grad_out):
         g.edata[grad_score_name] = out * grad_out
         g.update_all(fn.copy_e(grad_score_name, 'm'), fn.sum('m', accum_name))
         g.apply_edges(fn.e_mul_v(out_name, accum_name, out_name))
-        grad_score = g.edata[grad_score_name] - g.edata[out_name]
+        g.ndata.pop(accum_name)
+        grad_score = g.edata.pop(grad_score_name) - g.edata.pop(out_name)
         return None, grad_score