@enisberk
you’re the lifesaver !!!, thanks a lot
But I don’t understand how did you come up with the solution? Do you mind sharing how the code works?
The scenario can be described like this:
lookup=torch.ones((15, 720, 30))
means:  batch_size x num_nodes x feature_size
where num nodes are the number of nodes in a tree and each node is represented by an embedding with size = feature_size.
children= (15, 720, 19, 2)
means: batch_size x num_nodes x num_children x 2. I don’t really understand the meaning of 2 here.
Since every node in a tree is represented by an embedding, what the tensorflow code did is that they want to “merge” the 2 tensors into one, the dimension 4 means that each child in dimension 3 contains a corresponding embedding of its parent, make it more convenient for the later step.
In your code, i don’t really get how it works but it did the job…