I’m trying to implement a Linear layer with an extra constant weight weight_fb to compute grad_input in backward pass.
So I implemented a MyLinearFunc Function similar to the one here 1. In forward function of MyLinear Module, return
MyLinearFunc.apply(input, self.weight, self.weight_fb, self.bias)

and the return of backward function in MyLinearFunc is
grad_input, grad_weight, None, grad_bias

When I initiate weight_fb as
self.weight_fb = nn.Parameter(torch.Tensor(out_features, in_features), requires_grad=False)

everything works fine.
But as weight_fb is not a learnable parameter, I thought it would be better to use
self.weight_fb =torch.Tensor(out_features, in_features)
self.register_buffer('feedback_weight', self.weight_fb)

But this gives CUDA error: an illegal memory access was encountered when using GPU.
Just wondering if nn.Parameter and Tensor with register_buffer() give the same training results, and if I did something wrong when using register_buffer().
Any help would be appreciated!