Tree ops to CPU (#12)

qipeng · May 28, 2019 · db7c128 · db7c128
1 parent e574f4d
commit db7c128
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/model/gcn.py b/model/gcn.py
@@ -74,6 +74,7 @@ def forward(self, inputs):
         maxlen = max(l)
 
         def inputs_to_tree_reps(head, words, l, prune, subj_pos, obj_pos):
+            head, words, subj_pos, obj_pos = head.cpu().numpy(), words.cpu().numpy(), subj_pos.cpu().numpy(), obj_pos.cpu().numpy()
             trees = [head_to_tree(head[i], words[i], l[i], prune, subj_pos[i], obj_pos[i]) for i in range(len(l))]
             adj = [tree_to_adj(maxlen, tree, directed=False, self_loop=False).reshape(1, maxlen, maxlen) for tree in trees]
             adj = np.concatenate(adj, axis=0)