Thanks for your suggestion @albanD,
When I changed it to the following, it worked:
x = F.linear(x,nn.Parameter(torch.randn(self.feature_dimension, x.shape[1])))
Also, just a small question, so here the parameters are learnable right (asking in general about the parameters of torch.nn.functional)?
Thanks