To me, your implementation of custom conv seems to be correct, except two things in testing code:

In testing code, the two conv layers are not sharing the same weight. You can assign the weight of one conv layer to the other as follows:


    conv1 = nn.Conv2d(input.shape[1], 1, kernel_size=3, dilation=1, padding=1, stride=1 ,bias = False)
    conv2 = customConv(n_channels=input.shape[1], out_channels=1, kernel_size=3,  dilation=1, stride =1, padding = 1, bias = False)
    conv1.weight = conv2.weight
    conv1.bias = conv2.bias


Once you assign the same weights, the conv outputs are very close (maybe only varies in numerical precision). You can check using L2 loss:

print(torch.pow(conv1_output - conv2_output, 2).sum())

Maybe, take a simple 5 x 5 input and verify the output by printing.