Hello Everyone,
I am building a network with several graph convolutions involved in each layer.
A graph convolution requires a graph signal matrix X and an adjacency_matrix adj_mx
The network simplified computation graph looks as follow:

Screenshot 2020-04-19 at 11.51.391576×1196 33.2 KB

In (a) the network has self.adj_mx being used in all layers.
In (b) I added a learnable mask adj_mx_mask for the adj_mx.
We have therefore self.adj_mx_mask = nn.Parameter(…)  same shape as self.adj_mx
In each layer, self.adj_mx = self.adj_mx * self.adj_mx_mask
The problem is, with (a), loss.backward() works perfectly well.
With (b) I get the error to use loss.backward(retain_graph=True) which makes the training process way slower.
I need a way to add this learnable mask without having to make the training so slow.
Any idea on how I could proceed with that?
Thanks!