[BugFix] Fix performance bug in EdgeSoftmax (#452)

* fix performance bug in EdgeSoftmax

* minor

python/dgl/nn/pytorch/softmax.py

@@ -84,6 +84,8 @@ class EdgeSoftmax(nn.Module):
         graph.apply_edges(
             lambda edges: {self._logits_name : th.exp(edges.data[self._logits_name] -
                                                       edges.dst[self._max_logits_name])})
+        # pop out temporary feature _max_logits, otherwise get_ndata_name could have huge overhead
+        graph.ndata.pop(self._max_logits_name)
         # compute normalizer
         graph.update_all(fn.copy_edge(self._logits_name, self._logits_name),
                          fn.sum(self._logits_name, self._normalizer_name))