PyTorchで行列の各行に対してある列の値を取り除いた行列を得たい（高速化）

例として以下のような5x5行列を考えます。

Python
1import torch      
2x = torch.rand(25).reshape(5, 5)  
3x
4
5>>>
6tensor([[0.9263, 0.6601, 0.3334, 0.6175, 0.6035],
7        [0.2583, 0.6105, 0.3113, 0.2965, 0.9429],
8        [0.0350, 0.9206, 0.8667, 0.8958, 0.7814],
9        [0.8921, 0.8116, 0.7271, 0.1324, 0.8097],
10        [0.7166, 0.7780, 0.5185, 0.4530, 0.4059]])

また各行ごとにある列のindexが与えられているとします。

drop_indices = torch.tensor([1, 4, 0, 3, 2])          
drop_indices

>>>
tensor([1, 4, 0, 3, 2])

この時に、i行に対してdrop_indices[i]列の値を消したような5x4の行列を得たいです。
一行目なら一列目の0.6601、二行目なら四列目の0.9429をその行から除きたいということです。
つまり以下のような行列を結果として得たいです。

tensor([[0.9263, 0.3334, 0.6175, 0.6035],
        [0.2583, 0.6105, 0.3113, 0.2965],
        [0.9206, 0.8667, 0.8958, 0.7814],
        [0.8921, 0.8116, 0.7271, 0.8097],
        [0.7166, 0.7780, 0.4530, 0.4059]])

forで回してconcatで一応やりたいことはできるのですが、実際には1000行、5000列くらいの行列である程度高速に動いて欲しいです。（深層学習の学習処理の一部です）

どのような方向性で行けば良いでしょうか。どうぞよろしくお願いいたします。

行動規範の内容に同意します

回答1件

自己解決

以下のように自己解決しました。ありがとうございました。

Python
1>>> import torch
2>>> x = torch.rand(25).reshape(5, 5)  
3>>> drop_indices = torch.tensor([1, 4, 0, 3, 2])   
4>>> drop_elements = x[torch.tensor(range(5)), drop_indices].reshape(5, 1)
5>>> mask = x == drop_elements
6>>> x_2 = x[~mask].reshape(5, 4) 
7>>> x
8tensor([[0.3465, 0.5554, 0.9142, 0.1674, 0.1353],
9        [0.9194, 0.2897, 0.7397, 0.5700, 0.9786],
10        [0.5404, 0.5266, 0.4050, 0.4092, 0.1816],
11        [0.9258, 0.0706, 0.9894, 0.8694, 0.3407],
12        [0.8124, 0.0562, 0.1115, 0.4929, 0.2795]])
13>>> drop_indices
14tensor([1, 4, 0, 3, 2])
15>>> x_2
16tensor([[0.3465, 0.9142, 0.1674, 0.1353],
17        [0.9194, 0.2897, 0.7397, 0.5700],
18        [0.5266, 0.4050, 0.4092, 0.1816],
19        [0.9258, 0.0706, 0.9894, 0.3407],
20        [0.8124, 0.0562, 0.4929, 0.2795]])
21