def METHOD_NAME(node, diff_node):
    if node.kind() == "prim::DifferentiableGraph":
        diff_node.append(node)
    else:
        for block in node.blocks():
            for n in block.nodes():
                METHOD_NAME(n, diff_node)