Вы также можете просто назначить это конкретное измерение.
orig = torch.randint(low=0, high=10, size=(2,3,2,2))
fake = torch.randint(low=111, high=119, size=(2,1,2,2))
orig[:,[2],:,:] = fake
Оригинал до
tensor([[[[0, 1],
[8, 0]],
[[4, 9],
[6, 1]],
[[8, 2],
[7, 6]]],
[[[1, 1],
[8, 5]],
[[5, 0],
[8, 6]],
[[5, 5],
[2, 8]]]])
Подделка
tensor([[[[117, 115],
[114, 111]]],
[[[115, 115],
[118, 115]]]])
Оригинал после
tensor([[[[ 0, 1],
[ 8, 0]],
[[ 4, 9],
[ 6, 1]],
[[117, 115],
[114, 111]]],
[[[ 1, 1],
[ 8, 5]],
[[ 5, 0],
[ 8, 6]],
[[115, 115],
[118, 115]]]])
Надеюсь, это поможет!:)