PyTorch torch.index_select 関数
torch.index_selectPyTorch において、指定した次元に沿ってインデックスに対応する要素を選択するための関数です。
関数定義
torch.index_select(input, dim, index)
使用例
例
import torch
x = torch.randn(4, 5)
# 0行目と2行目を選択
indices = torch.tensor([0, 2])
result = torch.index_select(x, dim=0, index=indices)
print("結果の形状:", result.shape)
x = torch.randn(4, 5)
# 0行目と2行目を選択
indices = torch.tensor([0, 2])
result = torch.index_select(x, dim=0, index=indices)
print("結果の形状:", result.shape)
その他の拡張
Pytorch torch リファレンスマニュアル