After trial and error, I still could not find where the bug is
class MyBN1(nn.BatchNorm2d):
    def __init__(self, num_features, eps=1e-5, momentum=0.1,
                 affine=True, track_running_stats=True):
        super(MyBN1, self).__init__(
            num_features, eps, momentum, affine, track_running_stats)

    def forward(self, input):
        self._check_input_dim(input)

        exponential_average_factor = 0.0

        if self.training and self.track_running_stats:
            if self.num_batches_tracked is not None:
                self.num_batches_tracked += 1
                if self.momentum is None:  # use cumulative moving average
                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                else:  # use exponential moving average
                    exponential_average_factor = self.momentum

        # calculate running estimates
        if self.training:
            mean = input.mean([0, 2, 3])
            var = input.var([0, 2, 3], unbiased=False)
            n = input.numel() / input.size(1)
            with torch.no_grad():
                self.running_mean = exponential_average_factor * mean\
                    + (1 - exponential_average_factor) * self.running_mean
                # update running_var with unbiased var
                self.running_var = exponential_average_factor * var * n / (n - 1)\
                    + (1 - exponential_average_factor) * self.running_var
            mean = mean.view(1, -1, 1, 1)
            var = var.view(1, -1, 1, 1)
                
        else:
            mean = self.running_mean
            var = self.running_var
        weight = self.weight.view(1, -1, 1, 1)
        bias = self.bias.view(1, -1, 1, 1)

        output = MyBNFunc.apply(input, mean, var, weight, bias, self.eps)
        
        return output

def mysum(tensor):
    return tensor.sum((0,2,3),keepdim=True)
class MyBNFunc(Function):
    @staticmethod
    def forward(ctx, input, avg, var, gamma, beta, eps):
        ctx.avg = avg
        ctx.var = var
        ctx.eps = eps
        ctx.shape = input.shape
        output = input - avg
        scale = 1 / torch.sqrt(var + eps)
        output = output * scale
        ctx.save_for_backward(input, gamma, beta, output, scale)
        output = output * gamma + beta
        return output
    
    @staticmethod
    def backward(ctx, grad_output):
        input, gamma, beta, output, scale = ctx.saved_tensors
        avg = ctx.avg
        var = ctx.var
        eps = ctx.eps
        B, C, H, W = ctx.shape

        dL_dxi_hat = grad_output * gamma
        dl_dvar = mysum(dL_dxi_hat) * (input - avg) * -0.5 * scale * scale * scale
        dl_davg = mysum(dL_dxi_hat * -1.0 * scale) + dl_dvar * mysum(-2.0 * (input - avg)) / (B * H * W)
        dL_dxi = dL_dxi_hat * scale + dl_dvar * 2.0 * (input - avg) / (B*H*W)  + dl_davg / (B*H*W)
        dL_dgamma = (grad_output * output).sum((0, 2, 3), keepdim=True) 
        dL_dbeta = (grad_output).sum((0, 2, 3), keepdim=True)
        return dL_dxi, dl_davg, dl_dvar, dL_dgamma, dL_dbeta, None


compared to the result of torch.nn.BatchNorm2d, the gradient of gamma and beta match exactly, but the input_grad is wrong, could anyone help me?