So, how does it work when I have multiple channels?
For example, the below code will output tensors of shape (B, 2, W, H). Instead, I would like the two kernels to be each applied to each channel with output image (B, 2C, W, H)
def get_grad_kernel(channels):
  Iy = [[ 1, 2, 1],
        [ 0, 0, 0],
        [-1,-2,-1]]
  Ix = [[-1, 0, 1],
        [-2, 0, 2],
        [-1, 0, 1]]
  return torch.Tensor([ [Ix]*channels, [Iy]*channels ])/4

grad_kernel = get_grad_kernel(channels=3)
compute_gradient = lambda image: F.conv2d(input=image, weight=grad_kernel, padding=1)

EDIT: From ptrblck’s reply, making the kernel as the following
  return torch.Tensor([ [Ix], [Iy] ]*channels)/4

and setting groups=image.size(1) works nicely