[Pytorch select] index function

select(dim, index): The first parameter is the dimension of the index, and the second parameter is the serial number of the dimension of the index

import torch
a = torch.rand((3, 4))
print(a)
>>> tensor([[0.8664, 0.9759, 0.3063, 0.0686],
	        [0.6778, 0.0574, 0.3194, 0.4253],
	        [0.5045, 0.8318, 0.1745, 0.3150]])
print(a.select(dim=1, index=1)) # 取第1个维度中索引为1的值
>>> tensor([0.9759, 0.0574, 0.8318])

Use copy_ assignment to modify the tensor

import torch
a = torch.rand((3, 4)
print(a)
>>> tensor([[0.6115, 0.2551, 0.8714, 0.3236],
	        [0.3369, 0.4372, 0.2083, 0.4733],
	        [0.0046, 0.0981, 0.9148, 0.7852]])
y = torch.tensor([3, 2, 1])
a.select(1, 1).copy_(y.data) # 将第1个维度中索引为1的张量替换为y
print(a)
>>> tensor([[0.6115, 3.0000, 0.8714, 0.3236],
	        [0.3369, 2.0000, 0.2083, 0.4733],
	        [0.0046, 1.0000, 0.9148, 0.7852]])

Guess you like

Origin blog.csdn.net/weixin_43486780/article/details/111612420