If I want to add features for my edges, can I simply replace the encoder network with a network that supports edge_attr?

Hi Petar,
I have read your idea abot Deep Graph Infomax.
While right now something really puzzles me: if I want to add features for my edges and I still want to do unsupervised or self-supervised learning. Can I just simply change the encoder and remain all the other function unchanged? Such as corruption function or readout function?