Select values along one dimension

The question is if I have a 2-dimensional tensor [ [0, 1], [2, 3], [0, 0] ], how to write code without any loop to get the result [ [2], [1, 3] ] ? The zeros denote invalid positions which we do not select.

Hi,

I suppose that the expected result is [1, 2, 3] and not [[2], [1, 3]]. Here is one way you can do it:

    x = torch.tensor([[0, 1], [2, 3], [0, 0]], dtype=torch.int32)
    mask = (x != 0)
    y = x[mask]
    print(y)
    # tensor([1, 2, 3], dtype=torch.int32)

Let me know if it is not what you wanted!

Thanks for your answer. But it is not what I want. The result I want, [ [2], [1, 3] ], reflects that the element 2 is selected from the first column of [ [0, 1], [2, 3], [0, 0] ], and the elements 1 and 3 are selected from the second column of [ [0, 1], [2, 3], [0, 0] ].

Ok, the result you want is a tensor? Because I am pretty sure that you cannot have a tensor with a varying number of columns. Or do you want your result as Python lists?

Yes, a python list is ok.

maybe we should first convert [ [0, 1], [2, 3], [0, 0] ] to python list, because pytorch build-in functions always return tensors which must have the same number of numbers along one dim.

Hmm, the problem with converting to python lists is now that we have to use for loops… I am afraid that I am not able to help you with your problem, the only way I arrive at your expected result involves a for loop.

Sorry, and best of luck!

Thanks for your patience.