Hi @tom, yes, the connected components are disjoint.

and then indexing with the mask to get the values back into the index.

Could you elaborate this part a bit more please ?
Update: This is what I could come up from your approach, but it still has the same bottleneck for loop,
import torch
from torch_scatter import scatter_max

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.]]]])


indexes = torch.unique(CCs)
out = inp.new_zeros(indexes.shape)
out, argmax = scatter_max(inp.view(-1), CCs.view(-1), out=out)
maxified = inp.new_ones(inp.shape)

for i in range(len(indexes)):
    mask = indexes[i] == CCs
    maxified[mask.view((1, 1) + mask.shape)] = out[i]

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