Hi,
I am trying to run inference a custom network with following state_dict and corresponding weight and bias dimensions loaded from a pre-trained model. I have attached part of the network in which the error seems to be occurring.
.
.
.
name=encoder.original_model.classifier.weight, size=torch.Size([1000, 2208])
name=encoder.original_model.classifier.bias, size=torch.Size([1000])

name=decoder.conv2.weight, size=torch.Size([1104, 2208, 1, 1])
name=decoder.conv2.bias, size=torch.Size([1104])

name=decoder.up1.convA.weight, size=torch.Size([552, 1488, 3, 3])
name=decoder.up1.convA.bias, size=torch.Size([552])
name=decoder.up1.convB.weight, size=torch.Size([552, 552, 3, 3])
name=decoder.up1.convB.bias, size=torch.Size([552])
name=decoder.up2.convA.weight, size=torch.Size([276, 744, 3, 3])
name=decoder.up2.convA.bias, size=torch.Size([276])
name=decoder.up2.convB.weight, size=torch.Size([276, 276, 3, 3])
name=decoder.up2.convB.bias, size=torch.Size([276])
name=decoder.up3.convA.weight, size=torch.Size([138, 372, 3, 3])
name=decoder.up3.convA.bias, size=torch.Size([138])
name=decoder.up3.convB.weight, size=torch.Size([138, 138, 3, 3])
name=decoder.up3.convB.bias, size=torch.Size([138])
name=decoder.up4.convA.weight, size=torch.Size([69, 234, 3, 3])
name=decoder.up4.convA.bias, size=torch.Size([69])
name=decoder.up4.convB.weight, size=torch.Size([69, 69, 3, 3])
name=decoder.up4.convB.bias, size=torch.Size([69])
name=decoder.conv3.weight, size=torch.Size([1, 69, 3, 3])
name=decoder.conv3.bias, size=torch.Size([1])

I have implemented a custom convolution operation in the decoder.conv2 layer and the implementation code snippet is attached below.
class approx_conv(torch.autograd.Function):

	@staticmethod

	def forward(ctx, in_feature, kernel, out_channel, padding, bias):

		batch_size = in_feature.size(0)
		in_channels = in_feature.size(1)
		orig_h, orig_w = in_feature.size(2), in_feature.size(3)

		kh, kw = kernel.size(2), kernel.size(3)
		dh, dw = 1, 1
		pad = padding

		out_h = (orig_h+2*pad-kh)//dh + 1
		out_w = (orig_w+2*pad-kw)//dw + 1

		img = F.pad(input= in_feature, pad= (pad, pad, pad, pad), mode='constant', value= 0)

		h, w = img.size(2), img.size(3)
		patches = img.unfold(2,kh,dh).unfold(3,kw,dw).reshape(batch_size, in_channels*kh*kw, -1)
		result = torch.zeros(batch_size, out_channel, out_h*out_w)

		for b in range(batch_size):
			x = patches[b]
			for i in range(kernel.size(0)):
				start = time.time()
				for j in range(x.size(1)):
					r=0
					for k in range(kernel.size(1)):
					
						t1 = int((kernel[i][k]*1000).round())
						t2 = int((x[k][j]*1000).round())
						
						if(t1>255): t1 =255
						if(t1<-255): t1 =-255
						if(t2>255): t2=255
						if(t2<-255): t2=-255
						
						if((t1>0 and t2>0) or (t1<0 and t2<0)):
							t1=abs(t1)
							t2=abs(t2)
							sign=1
						else:
							t1=abs(t1)
							t2=abs(t2)
							sign=-1
							
						r+= lookup_table[t1][t2]*sign
						
					result[b][i][j] = r/1000000
				end = time.time()
				print("time_taken for 300 conv in {} / 1104 = {}".format(i,end-start))
				
		result = result.reshape(batch_size, out_channel, out_h, out_w)

		if bias is not None:
			result += bias.unsqueeze(0).expand_as(result)
		
		return result

In this case bias.shape = torch.Size([1104])  and result.shape = torch.Size([1,1104,15,20])
I am getting the error in the line: result += bias.unsqueeze(0).expand_as(result)
Traceback (most recent call last):
  File "evaluate_py.py", line 47, in <module>
    e = evaluate(ap_model, rgb, depth, crop, batch_size=1)
  File "/gdrive/My Drive/DenseDepth_DenseNet_161/pytorch_densenet161/utils2.py", line 71, in evaluate
    pred_y = scale_up(2, predict(model, x / 255, minDepth=10, maxDepth=1000, batch_size=bs)[:, :, :, 0]) * 10.0
  File "/gdrive/My Drive/DenseDepth_DenseNet_161/pytorch_densenet161/utils2.py", line 35, in predict
    predictions = model(images)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/gdrive/My Drive/DenseDepth_DenseNet_161/pytorch_densenet161/approx_Model.py", line 61, in forward
    return self.decoder( self.encoder(x) )
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/gdrive/My Drive/DenseDepth_DenseNet_161/pytorch_densenet161/approx_Model.py", line 36, in forward
    x_d0 = self.conv2(x_block4)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/gdrive/My Drive/DenseDepth_DenseNet_161/pytorch_densenet161/customconv_v2.py", line 129, in forward
    return approx_conv.apply(x, self.weight, self.out_channels, self.padding, self.bias)
  File "/gdrive/My Drive/DenseDepth_DenseNet_161/pytorch_densenet161/customconv_v2.py", line 106, in forward
    result += bias.unsqueeze(0).expand_as(result)
RuntimeError: The expanded size of the tensor (20) must match the existing size (1104) at non-singleton dimension 3.  Target sizes: [1, 1104, 15, 20].  Tensor sizes: [1, 1104]

I have implemented the forward() function just as suggested in the Pytorch documentation for making custom layers. Is there a workaround for this error?
When I reshape the result tensor in the following manner I do not get errors:
result = result.reshape(batch_size, out_h, out_w, out_channels)
But is the right way to go ahead?
TIA