Solved by tueboesen in post #2 
              
              
                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)
    …