Two quick comments:

Use max(inp.view(B, C, -1), 2) and numpy.unravel_index (or just // and %) after.
You cannot differentiate argmax, it’s mathematically impossible. People suggested various things (e.g. with softargmax or argsoftmax - I cannot remember - that was defined as torch.arange(W*H).view(1, 1, W * H)*torch.softmax(inp.view(B, C, -1), dim=2) or somesuch. It’s not always entirely clear how to interpret that, but it seems that some people like it.

Best regards
Thomas