pytorchでTensorで一番値が大きいインデックスの場所以外の値を0にしたい

Question

Pytorchで1次元のtensorがあり、その中で最も値が高いインデックスをargmaxで求めました。
ここで、tensorの中で値が一番大きい場所（ここではargmaxで求めた2番目）以外の値を0にしたいのですがどうすればいいでしょうか？
2番目の値はそのままにしたいです。

ご教授よろしくお願いします。

inp1 = torch.rand((1, 5), dtype=torch.float32)
idx = torch.argmax(inp1)
print(inp1, idx)
tensor([[0.6886, 0.0348, 0.9323, 0.4553, 0.6799]]) tensor(2)

以下のようにしたい
tensor([[0, 0, 0.9323, 0, 0]])

Answer 1

numpyもそうですが、元配列に対するbooleanな結果を出す式を「マスク配列」と呼びます。
例えば[True, True, False]みたいな配列がマスク配列です。

元の配列[マスク配列] = 一致したときに置き換えたい値
という構造でやりたいことができます。

参考：　この記事の「リスト1.3　Tensorのインデクシング操作例」を読んでください。

以下、走らせてませんがたぶん動きます。

inp1 = torch.rand((1, 5), dtype=torch.float32)
idx = torch.argmax(inp1)
print(inp1, idx)
# tensor([[0.6886, 0.0348, 0.9323, 0.4553, 0.6799]]) tensor(2)

#########################
# 以下、追加コード
inp1[inp1 != torch.max(inp1 )] = 0

---

例えば、
生き物 = [くじら, シャチ, サメ, ペンギン]
生き物[生き物==魚類] = None # 魚類だけNoneに置き換える
print(生き物)
[くじら, シャチ, None, ペンギン] # サメがNoneに置換された！

みたいになります。

Comments

[trafalbad]

解決しました。ありがとうございます！