Thanks @ptrblck, your suggestion helped me very much.
I added a few markers, so now I can tell which parts of the code the memory operations refer to.
I also start the profiler after the first few itereations, as you suggest.
I post here the output I currently get from NVVP, it might be useful for other people in the future.
nvvp_both_new2451×1010 138 KB
I added markers for the edge update and the node update of my GNN, plus one marker for the loss function.
The loss function is clearly taking a lot more time in the new version of the code, so I will start my code optimisation from there.