Create a new layer with transpose matrix

1

Hi, I have to define a new layer, which does the following:

Given an input x (vector with N elements), and a NxM matrix W, I want as an output W^T(ReLU(Wx)), where W^T is the transpose of W. What I have is the following:

class AML(nn.Module):
    def __init__(self, in_features: int, out_features: int,n: int) -> None:
        super(AML, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        param = torch.Tensor(out_features, in_features)
        param_t = torch.transpose(param,0,1)
        self.weight = Parameter(param)
        self.weight_t = Parameter(param_t)
        self.relu = ReLu(n)

    def forward(self, input: Tensor) -> Tensor:
        out = F.linear(input, self.weight)
        out = self.relu(out)
        out = F.linear(out,self.weight_t)
        return out

But I keep getting random guessing results even on very simple tasks. Is there anything wrong is my code?

2

The code looks generally alright, but I would not use torch.Tensor to create the initial parameter, as it will use uninitialized memory and will contain whatever was stored in the memory you are using.
This would also mean that it might contain NaN values etc.
Use the factory methods, such as torch.randn etc. instead.