At the moment, I’m experimenting with defining custom sparse connections between two fully connected layers of a neural network.

To accomplish this, right now I’m modifying  nn.Linear(in_features, out_features) to nn.MaskedLinear(in_features, out_features, mask), where mask is the adjacency matrix of the graph containing the two layers. The module nn.Linear uses a method invoked as self._backend.Linear() defined in nn._functions.linear.py, which I clearly will have to modify as well.
I’m not sure if this is the right way to do it. I have a constant feeling that there should be a better way to do this. Any help/comments on this are are much appreciated,