Pytorch 使用索引张量选择张量的第二个维度

我有一个2D张量和一个索引张量。2D张量有一个批次维度和一个包含3个值的维度。我有一个索引张量，它正好选择了3个值中的1个元素。生成仅包含索引张量中元素的切片的“最佳”方法是什么

t = torch.tensor([[1,2,3], [4,5,6], [7,8,9]])
t = tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])

i = torch.tensor([0,0,1], dtype=torch.int64)
tensor([0, 0, 1])

预期输出

tensor([1, 4, 8])

---

答案的一个例子如下

导入火炬
t=火炬张量（[[1,2,3]，[4,5,6]，[7,8,9]）
col_i=[0,0,1]
第i行=范围（3）
打印（t[第一行，第一列]）
#张量（[1,4,8]）

答案的示例如下

导入火炬
t=火炬张量（[[1,2,3]，[4,5,6]，[7,8,9]）
col_i=[0,0,1]
第i行=范围（3）
打印（t[第一行，第一列]）
#张量（[1,4,8]）