How to use where function like Numpy

[Ryohei_Waitforit_Iwa (Ryohei Waitforit Iwata)]

I ran the following code, got an error.

@Ryohei_Waitforit_Iwa

a = torch.Tensor([[1, 2, 3], [4, 5, 6]])
torch.where(a[:, 0] == 1)

TypeError: where() missing 2 required positional argument: “input”, “other”

Numpy allows us to use where function without 2 arguments, but Pytorch not.
What I want to do is to select the rows corresponding to the value what I put.
I will really appreciate that if you would told me how to use this function.

[albanD]

From looking at numpy’s doc when only a single argument is given, it is equivalent to condition.nonzero(). So just do (a[:, 0] == 1).nonzero() ?