I am new to pytorch and I want to write a layer in which the input/output connecteion
is defined through an adjacency matrix, and all existing edges should
have same weight. My current implementation is as follows:
class myFunc(Function):
    def __init__(self):
        super(myFunc, self).__init__()
    def forward(self, input, weight,adj):
        a = torch.zeros_like(adj,dtype=torch.float)
        a[adj] = weight
        return input.mm(a)

class MyLayer(nn.Module):
    def __init__(self,adj):
        super(MyLayer, self).__init__()
        self.adj = adj
        self.weight = nn.Parameter(data=torch.Tensor(1), requires_grad=True)
        nn.init.normal_(self.weight.data, mean=1, std=0.1)

    def forward(self, input):
        return myFunc().forward(input,self.weight,self.adj)

This approach does not use less memory than the case that I have a separate weight per net.
Since I have a single parameter to learn (one weight for the entire layer), is there any way that I can implement this in a more efficient way?