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?