Invalid output is provided when input is not contiguous

[botcs]

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