Also I am not sure how you know at the end that  `nodeVecs[root]`  actually contains what you want.
My hope is that information from the entire tree gets bubbled up to the root node over the iterations and that the root would then become a suitable representation of the entire tree. I was hoping that backpropogation would be able to learn this. I could also do max-pooling or averaging over the nodes, but I figured for a tree, this might work better?
I think I see what you mean with the children discovery being moved out. This would make it two doubly nested for loops instead of one triply nested for loop, correct? (I’d need two for loops to do the children discovery since I need to find all children of all nodes.) I’ll try that and see if it gets much better.
Thanks for the suggestion!
More generally, do you have any tips for speeding up code which isn’t easily vectorized into pytorch and which might generate large computational graphs for which backprop might be slow?