技術標籤：pytorch

select(dim, 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])

使用copy_賦值修改張量

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]])