Unfortunately, I don’t understand how your output tensor is created given the inp and masks (e.g. where is the 11. coming from in output[0, 0, 0, 2]).

Sorry for the poor explanation. Each value in CCs indicates which connected component that pixel belongs to. So in the above example, 3 pixels belong to 0th CC, 1 pixel belongs to 1st CC and the rest belong to 2nd CC.
For custom pooling, all pixels of a CC are set to their maximum value from the input image. These are 11. for the 0th CC, 12. for the 1st CC and again 12. for the 2nd CC.
The problem with the summed mask approach is that each pixel exists in a connected component. Therefore, masks.sum(dim=0).bool() results in a mask filled only with True values, and applying it to input will return input again. I added the operation, the for loop in forward is causing a bottleneck.
import torch

class AdvMaxify(torch.autograd.Function):
    
    @staticmethod
    def forward(ctx, inp, CCs):
        maxified = torch.ones(inp.shape, device=inp.device)
        cs = torch.unique(CCs)
        mask_shape = cs.shape + maxified.shape
        masks = CCs.expand(mask_shape) == cs.view(cs.shape + (1, 1, 1, 1)).expand(mask_shape)
        for m in torch.unbind(masks, dim=0): # bottleneck
            m_inp = inp * m
            dim_maxes = torch.amax(m_inp.view(inp.shape[0], inp.shape[1], -1).squeeze(), dim=-1)
            num_dims = m.dim() - dim_maxes.dim()
            new_shape = dim_maxes.shape + (1,) * num_dims if m.shape[0] > 1 else (1,) + dim_maxes.shape + (1,) * (num_dims - 1)
            maxified[m] = dim_maxes.view(new_shape).expand_as(m)[m]
        return maxified
    
    @staticmethod
    def backward(ctx, grad_output):
        return grad_output, None


if __name__ == "__main__":

    inp = torch.tensor([[[[11., 11., 10., 12.],
                          [10., 11., 11., 10.],
                          [11., 12., 12., 12.],
                          [10., 12., 11., 10.]]]], requires_grad=True)

    CCs = torch.tensor([[0., 0., 0., 1.],
                        [2., 2., 2., 2.],
                        [2., 2., 2., 2.],
                        [2., 2., 2., 2.]])

    reference_output = torch.tensor([[[[11., 11., 11., 12.],
	                                   [12., 12., 12., 12.],
	                                   [12., 12., 12., 12.],
	                                   [12., 12., 12., 12.]]]])

    adv_maxify = AdvMaxify.apply
    outp = adv_maxify(inp, CCs)

print((outp == reference_output).all())
> tensor(True)