How to make weight multiple a Parameter before forward?

I want to multiple the weight by a Parameter beta, and after that, the Parameter beta can be update bu Pytorch, How can I do?

here is my code, but it doesn’t work.

class BinaryWeight(torch.autograd.Function):
def init(self):
self.u = nn.Parameter(torch.Tensor([0.0]), requires_grad=True)
self.alpha = nn.Parameter(torch.Tensor([1.0]), requires_grad=True)

def forward(self, input):
    neg_mean = input.mean(dim=0, keepdim=True).neg()
    std = input.std(dim=0, keepdim=True)
    bw = input.add(neg_mean).add(self.u*std).sign().mul(self.alpha)
    return bw

 def backward(self, grad_output):
    return grad_output

input is conv’s weight
when module forward, I do such thing:
such as
conv1.weight = BinaryWeight()(conv1.weight)

I want to realize the function as below, how can I do?

Use the functional form, torch.nn.functional.conv_2d

Thank u, I will have a try!