Generally it is quite close to what I had in mind.



 Saikumar_Tadi:

1.should the_mask contain a tensor of 0’s and 1’s. 1’s at places with variables we want to compute the gradients for and 0’s in the places which we want to ignore gradient computation?


I’d make it a boolean tensor, so True for  trained False for fixed. (Mind you, see how neatly the first letters align? Programming is poetry. Just some form that is modern enough to be entirely tedious.)
The self.weight_fixed should be a buffer, too.


2.what is the significance of the self.weight_param,self.weight_fixed variables ? are these automatically computed by pytorch or should i define and initialise them in my init function ? if so how do i initialise them ?


self.weight_param should be a nn.Parameter and initialized like you would initialize plain parameters. Note that the part where the mask is False the parameter is ignored. The weight_fixed should be whatever you see fit. Here, the bits where the mask is True is ignored.