# A bug in torch.unique with return_inverse=True and dim != None?

Try this:

aa = torch.tensor([[1, 2, 2], [3, 3, 3]])
unq, groups = aa.unique_consecutive(dim=1, return_inverse=True)

groups is:
tensor([0, 1, 1])

which are the group assignments for the first row only! expected groups to be of shape [2, 3], with the group assignments for each row… Is this a bug? Any way of achieving what I intended in an efficient way?

Hi,

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

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