To add learnable parameters to an nn.Module, you should use nn.Parameter type:
self.W = nn.Parameter(torch.randn(dim, 1).type(dtype))
Then you need to make sure that this is passed to your optimizer with the other parameters of your network:
optimizer = optim.SGD(model.parameters() + criterion.parameters(), optim_args).
And no you don’t need to implement the backward pass if you do that