Hi,
You are reducing your function over second dimension, so it has to be in shape of [3, ] in your case.



 cruvadom:

dim=1


Also, it was just a coincidence that the returned groups is same as groups for just the first row!
When you set the dim=1, your function will be applied on given dimension rather than whole tensor as a flattened array. So, set dim=None.
You can run the following code in 4 cases,

with aa=aa1 and dim=None

With aa=aa1 and dim=1

with aa=aa2 and dim=None

With aa=aa2 and dim=1


aa1 = torch.tensor([[1, 0, 0],
                   [0, 0, 1]])
aa2 = torch.tensor([[1, 1, 1],
                   [0, 1, 1]])
aa = aa1  # change this to aa2 with dim=1 and dim=None and see the difference
print(aa)
unq, groups = aa.unique_consecutive(dim=1, return_inverse=True)
print(groups)
unq, groups = aa.unique_consecutive(dim=None, return_inverse=True)
print(groups)

Bests