Describe the bug
According to documentation this function should work for images with arbitrary number of channels, but current implementation fails for C<> 1.

Also documentation says that expected kernel shape is [n, k, k] while in reality we always pass [n, 1, k, k]. [n, k, k] is actually not supported due to the same error in implementation

To Reproduce

x = torch.randn(4, 3, 16, 16)
kernels = torch.randn(2, 1, 3, 3) # should be torch.randn(2, 3, 3) 

output = gradient_map(x, kernels)

Expected behavior
No errors, proper computation.

Additional context
Not sure yet, but probably this function also doesn't support kN != 2. Will check while fixing
UPD: Supports, but probably math for kN > 2 should be different. Won't dig into that now