Thanks a lot but there are 2 problems:
weight = nn.Parameter(torch.randn(2, 2))
optimizer = optim.SGD(weight, lr=0.001, momentum=0.9)



I received an error as “TypeError: params argument given to the optimizer should be an iterable of Tensors or dicts, but got torch.FloatTensor”


If I want to have a next layer (for example a linear layer in standard pytorch) after my custom layer, we have 2 sets of parameters: parameters related to next layer and parameters related to my custom layer. How could I merge them together and pass them to optimizer.


Many thanks