@albanD  I have a follow up question.
If i want to keep K1 as it is, but change K2 during training (with backpropagation), how should i do that?
K1 = torch.FloatTensor([[1,0,0],[0,-1,0],[0,0,0]]).unsqueeze(0).unsqueeze(0).to(device)
pad_k1 = 1
K2 = torch.FloatTensor([[0,1,0],[-1,0,0],[0,0,0]]).unsqueeze(0).unsqueeze(0).to(device)
pad_k2 = 1
def EdgeDetectionNew(Input_Temp):
   # Move the input to the right device (remove this line if this is not needed)
   Input_Temp = Input_Temp.to(device)
   # For the conv, make the channel look like batch
   # Compute the size for the conv
   inp_size = Input_Temp.size()
   conv_size = list(inp_size)
   conv_size[0] = inp_size[0] * inp_size[1]
   conv_size[1] = 1
   # View the input acordingly
   conv_input = Input_Temp.view(conv_size)
   Gx = F.conv2d(conv_input, K1, padding=pad_k1)
   Gy = F.conv2d(conv_input, K2, padding=pad_k2)
   GM = (Gx**2 + Gy**2)**0.5
   # View the output back to get the channel
   GM = GM.view(inp_size)

   return GM