tensor配列の要素のある場所を取得したいです．


tensor型の配列のインデックスを取得したいのですが，うまくいかず，教えて頂きたいです．以下のような形になっており，最大要素1がある場所だけを抽出したいです．

[code]

    tensor([[1., 0., 0., 0., 0., 0.],  →この場合は0を
        [1., 0., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0.],
        ...,
        [0., 0., 1., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0.]])　　→この場合は2を
    
[/code]

コード  

[code]

    import numpy as np
    import torch                                                                                   
    npzfile = np.load("C:/Users/namae/Desktop/myo-python-1.0.4/myo-armband-nn-master/data/train_set.npz")
    x = npzfile['x']
    y = npzfile['y']                                                                                     
    x = x.astype(np.float32)
    x = torch.from_numpy(x).clone()
    x = torch.tensor(x)
    x = x.clone().detach()
    y = y.astype(np.float32)
    y = torch.from_numpy(y).clone()
    print(y)
    
[/code]



[code]

    x = torch.tensor([[1., 0., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0.]])
    
    x.max(dim=1)[1]
    # tensor([0, 0, 0, 2, 2, 2])
    
[/code]

