tensor的一个索引方法

y torch.tensor([0, 2]) y_hat torch.tensor([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]]) y_hat[[0, 1], y]