pytorchで条件を満たす列だけ抽出したい

Question

閲覧ありがとうございます。
タイトルの通り、条件を満たすテンソルを抽出したいと考えております。
具体的には３行めが"0"の列を抽出したいです。
普通の配列なら、リスト内包表記を用いて下記の様に抽出できます。

>>> li=[[1, 2, 0], [3, 4, 1], [5, 6, 0], [7, 8, 0]]
>>> li = [row for row in li if row[2] == 0]
>>> li
[[1, 2, 0], [5, 6, 0], [7, 8, 0]]

---

しかしながら、同様の機能をpytorchのtensorに対して実行することができません。
tolistメソッドを用いてもリストにならず、無理やり内包表記を使うとリスト内にtensorが３つ入ってしまいます。

>>> import torch

>>> a = torch.tensor([[1, 2, 0], [3, 4, 1], [5, 6, 0], [7, 8, 0]])
>>> a
tensor([[1, 2, 0],
        [3, 4, 1],
        [5, 6, 0],
        [7, 8, 0]])

>>> a.tolist()
[[1, 2, 0], [3, 4, 1], [5, 6, 0], [7, 8, 0]]
>>> a #リストにならず？
tensor([[1, 2, 0],
        [3, 4, 1],
        [5, 6, 0],
        [7, 8, 0]])

>>> a = [row for row in a if row[2] == 0]
>>> a　#リスト内にtensorが３つ入ってしまった
[tensor([1, 2, 0]), tensor([5, 6, 0]), tensor([7, 8, 0])]

>>> a = torch.tensor(a) #無理やりtensorに戻そうとしてもうまくいかず...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: only integer tensors of a single element can be converted to an index

どなたかtensorに対して任意の行を抽出する方法をご存知でしたら、お教えください。

Answer 1

>>> a
tensor([[1, 2, 0],
        [3, 4, 1],
        [5, 6, 0],
        [7, 8, 0]])


>>> a[a[:, 2] == 0]
tensor([[1, 2, 0],
        [5, 6, 0],
        [7, 8, 0]])

Comments

[Moririn_1881]

まさに探していたものです。
ありがとうございます????

Answer 2

すごく簡単な回答で申し訳ないのですが（おそらく質問者さんもうっかり見落としているだけと思いますが…）
、tolist()は非破壊メソッドなので、
a = a.tolist()
とすれば期待通りの動きになると思います。
tensorのまま処理する方法については、後で思いついたら追記しておきます。

Comments

[Moririn_1881]

tolistは非破壊メソッドだったんですね、うっかりしてました。

ありがとうございます。