fix bug in treelstm

SunYa0 · Oct 8, 2018 · 1f0f57b · 1f0f57b
1 parent cdf7334
commit 1f0f57b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/examples/pytorch/tree_lstm/train.py b/examples/pytorch/tree_lstm/train.py
@@ -91,7 +91,7 @@ def _batcher(trees):
                 mean_dur = np.mean(dur)
                 print("Epoch {:05d} | Step {:05d} | Loss {:.4f} | "
                       "Acc {:.4f} | Time(s) {:.4f} | Trees/s {:.4f}".format(
-                    epoch, step, loss.item(), acc.item() / args.batch_size,
+                    epoch, step, loss.item(), acc.item() / len(label),
                     mean_dur, args.batch_size / mean_dur))
         print("Epoch time(s):", time.time() - t_epoch)
 

diff --git a/examples/pytorch/tree_lstm/tree_lstm.py b/examples/pytorch/tree_lstm/tree_lstm.py
@@ -23,7 +23,7 @@ def __init__(self, x_size, h_size):
         self.ut = 0.
 
     def message_func(self, src, edge):
-        return src
+        return {'h' : src['h'], 'c' : src['c']}
 
     def reduce_func(self, node, msgs):
         # equation (2)
@@ -90,9 +90,9 @@ def forward(self, graph, zero_initializer, h=None, c=None, iterator=None, train=
         """
         g = graph
         n = g.number_of_nodes()
-        g.register_message_func(self.cell.message_func, batchable=True)
-        g.register_reduce_func(self.cell.reduce_func, batchable=True)
-        g.register_apply_node_func(self.cell.apply_func, batchable=True)
+        g.register_message_func(self.cell.message_func)
+        g.register_reduce_func(self.cell.reduce_func)
+        g.register_apply_node_func(self.cell.apply_func)
         # feed embedding
         wordid = g.pop_n_repr('x')
         mask = (wordid != dgl.data.SST.PAD_WORD)