Backward() in custom conv2d

Hi everyone,

I am a beginner to pytorch and trying to build a custom convolution layer (conv2d) where instead of the default multiplication I replace it with my own multiplication algo. I’ve attached my implementation below.

I am getting the following error during backward pass - Function MBM_conv2dBackward returned an invalid gradient at index 1 - got [16, 32, 28, 28] but expected shape compatible with [32, 1, 3, 3]:

There seems to be a problem with the implementation of grad_weight calculation in the file:-

My doubts are the following:

  1. Where does needs_input_grad variable come from?
  2. input dimensions → [16,1,28,28] - (batch size = 16; 28x28 images)
    grad_weight dimensions → [32,1,3,3] - (32 3x3 kernels)
    What is the mistake being done during the grad_weight calculation? Or in the backward() implementation in general?

I couldn’t figure out the problem. Would really appreciate any help! Thanks!