If the 2D code is applied on a tensor that is a 2D slice of a larger tensor, e.g.
x.shape == [B, H, W, C]
correct_out = connected_components_labeling(x[0, :, :, 0].contiguous())
incorrect_out = connected_components_labeling(x[0, :, :, 0])
the expected output correct_out:

then incorrect_out will be a weird strided version of the correct_out:

the input was generated as follows, reusing the example code:
x = torch.tensor(cleared_copy).to("cuda", torch.uint8)[None, :, :, None].repeat(1024, 1, 1, 2)
Fixes
The quick fix is to call .contiguous() on every slice, the correct fix would be to either throw an error if the input is not contiguous or take the tensor strides into consideration on the cuda backend.
Hope this might help someone who got unexpected results.
Nevertheless, this implementations is very quick and is an awesome starting point, so thank you! <3
If the 2D code is applied on a tensor that is a 2D slice of a larger tensor, e.g.
the expected output

correct_out:then

incorrect_outwill be a weird strided version of thecorrect_out:the input was generated as follows, reusing the example code:
x = torch.tensor(cleared_copy).to("cuda", torch.uint8)[None, :, :, None].repeat(1024, 1, 1, 2)Fixes
The quick fix is to call
.contiguous()on every slice, the correct fix would be to either throw an error if the input is not contiguous or take the tensor strides into consideration on the cuda backend.Hope this might help someone who got unexpected results.
Nevertheless, this implementations is very quick and is an awesome starting point, so thank you! <3