Hello Everyone.
I am trying to implement the batch normalization layer and I am getting an error when trying to train using this layer.
The forward pass is correct, I am assuming because I tried to use a model containing the layer to get inference and it worked but for the backward I am having no idea what I should do exactly.
class NormGrad(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, alpha, betha, batch_size):
        ctx.save_for_backward(input, alpha, betha)
        input = input*alpha + betha
        return input

    @staticmethod
    def backward(ctx, grad_output):
        #still not finished
        input, alpha, betha = ctx.saved_tensors
        grad_input = None
        grad_betha = None
        grad_alpha = None
        #print(grad_input.size())
        grad_input= grad_output*alpha
        grad_alpha = grad_output.mul(input)
        grad_alpha = grad_alpha.sum(axis=(0,1,2))
        grad_bias = grad_output.sum(axis=(0,1,2))
        return grad_input, grad_alpha, grad_betha

So, I do not know whether what I am doing is correct or not but the first error that I am facing right now is saying that 4 parameters were expected to be returned but I am returning 3.