sahar.mi:

but without concatenation everything seems fine


My guess would be that you are creating a non-leaf tensor with the concatenation.
Could you check the is_leaf attribute of the parameters before and after the concatenation?
If it’s showing False afterwards (and True before), re-wrap the new tensor in an nn.Parameter.