that code was purely illustrative, register_buffer is just a way to store non-trainable tensors in a module (sometimes it is more conventient even for scalars).



 silvester:

How I could make a function in the way that it returns the gradients with respect to alphas, separately and not coupled with the other weights?


Why you want to manually return gradients in the first place? You usually only do this with external (thus non-differentiable) calculations, e.g. c++ extensions, or if you have simplified/more efficient gradient formulas. In other cases you only write forward() for nn.Module and most operations on tensors know how to calculate their gradients.