

Pytorchで3次元のtensorがあり、tensor中のdim=1の中で値が一番大きい場所を1に、一番大きい値以外を0にしたいのですがどうすればいいでしょうか？  
ご教授よろしくお願いします。

torch.argmaxやtorch.max、スライス処理など色々考えてみたんですがわからなかったです。

### 該当のソースコード

[code]

    import torch
    x = torch.randn(4, 3, 1)
    print(x)
    tensor([[[0.4082],
             [2.0627],
             [0.7252]],
    
            [[0.7946],
             [0.2679],
             [-0.4184]],
    
            [[0.3380],
             [0.8403],
             [-1.7227]],
    
            [[-1.1250],
             [-1.8144],
             [1.4441]]])
    
[/code]

### 以下のようにしたい

[code]

    print(x)
    tensor([[[0],
             [1],
             [0]],
    
            [[1],
             [0],
             [0]],
    
            [[0],
             [1],
             [0]],
    
            [[0],
             [0],
             [1]]])
    
[/code]


CC BY-SA 4.0
2023-02-01T05:49:28.480
93653
2023-02-03T07:01:32.850
2023-02-03T07:01:32.850
3060
56875
question
python 機械学習 pytorch
pytorchでTensorで一番値が大きいインデックスの場所以外の値を0にしたい


御質問へのコメントに書かれている他サイトの回答を参考にさせていただきました。  
（追記）コメントにて Python
のグローバルインタプリタロックの影響を受けるので，記述例の中のリスト内包表記は使わない方が無難との御指摘をいただきましたので修正しました。どのような影響があるのかと別解もコメント欄に記載していただきましたので参照ください。

[code]

    import torch
    
    torch.manual_seed(17)
    x = torch.randn(4, 3, 1)
    # y = torch.stack([(a == a.max()).int() for a in x])
    y = (x == x.amax(dim=1, keepdim=True)).int()
    
    print(f'{x}\n\n{y}')
    
[/code]

[code]

    tensor([[[-1.4135],
             [ 0.2336],
             [ 0.0340]],
    
            [[ 0.3499],
             [-0.0145],
             [-0.6124]],
    
            [[-1.1835],
             [-1.4831],
             [ 1.8004]],
    
            [[ 0.0096],
             [ 0.1534],
             [-2.6631]]])
    
    tensor([[[0],
             [1],
             [0]],
    
            [[1],
             [0],
             [0]],
    
            [[0],
             [0],
             [1]],
    
            [[0],
             [1],
             [0]]], dtype=torch.int32)
    
[/code]


CC BY-SA 4.0
2023-02-02T07:11:41.837
93680
2023-02-03T05:01:55.180
2023-02-03T05:01:55.180
54588
54588
93653
answer