I figured out the issue, the issue is that the print statement needs to be outside both with statements, like this:
    with profiler.profile(record_shapes=True) as prof:
        with profiler.record_function("model_inference"):
            node_vec = model(input=coords_init,xn_attr=node_attr)
    print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

It would be nice if this was included in the documentation/tutorial.