Is it possible to use norm.MessageNorm in a non-custom message passing layer?

[davidireland3]

Suppose you have a GNN that consists of some GCNConv (or any other built in) layers. Is there a way to use the MessageNorm to normalise the messages that are calculated internally to the method? I mean without manually adapting the source code yourself.

[rusty1s]

I think you should be able to just do this post GNN layer execution:

conv = GCNConv(16, 16)
norm = MessageNorm()

out = conv(x, edge_index)
out = norm(x, out)