For example
model = nn.Linear()
model.parameters() return a tuple of weight and bias. If I want to split the weight into several chunks along dimension 0. I can use
list(model.parameters())[0][:chunk_size,]
or
list(torch.split(param[0],chunk_size,0))[0]
But both of them will return a non-leaf node which cannot be the input of an optimizer.
One way to resolve is to write my own Linear() module and make weight isolated in advance. Then cat them together into self.weight and finally do things like
F.linear(input, self.weight, self.bias)

My question is whether I can just do some magic to the parameters of nn.Linear() that pytorch provided and achieve this goal. The reason is that it is tedious to write every module (nn.Con2d etc.) once again.