I thought some more about your suggestion of moving the children discovery out and I now don’t think it’s possible.
I’ve rewritten the code to hopefully be a bit more readable:
for i in range(getTreeDepth(tree)):
    fs = self.nodeEmbedding(nodeVecs)
    newNodeVecs = []
    for j in range(len(treeSymbols)):
        children = [fs[k] for k in nodeChildren[j]]                   
        newNodeVecs.append(self.neighborhoodEmbedding(fs[j], children))
    nodeVecs = torch.stack(newNodeVecs)

I don’t think I can move the “children discovery” out because it needs to reference the current vector values at the children. This depends on the outermost loop here. Do you see something I’m missing here?
Thanks,
Jack