如何将带梯度信息的张量复制替换到另一张量并保留梯度？

方法1：使用torch.index_put创建新张量 #

这是最简洁的实现方式，通过index_put生成新张量，完全避免in-place操作，同时完整保留计算图的梯度传递路径。

示例代码：

import torch

# 初始化示例张量
A = torch.randn(1, 768, requires_grad=True)
B = torch.randn(2, 4, 768, requires_grad=True)

batch = 0
replace_ids = [1, 3]  # 指定要替换的B中第batch个样本的位置

# 构造索引元组，对应B的前两个维度（batch维度、seq维度）
indices = (torch.tensor([batch]), torch.tensor(replace_ids))
# 生成新张量，accumulate=False表示直接替换目标位置的值而非累加
new_B = B.index_put(indices, A.repeat(len(replace_ids), 1), accumulate=False)

# 测试梯度传递
loss = new_B.sum()
loss.backward()
print(A.grad is not None)  # 输出True，A的梯度被正常保留
print(B.grad is not None)  # 输出True，原B未被修改，梯度正常传递

方法2：通过索引赋值创建新张量 #

如果需要更灵活的位置控制，可以手动克隆原张量后，仅对目标位置进行赋值操作。

示例代码：

import torch

A = torch.randn(1, 768, requires_grad=True)
B = torch.randn(2, 4, 768, requires_grad=True)

batch = 0
replace_ids = [1, 3]

# 克隆原B得到新张量，避免直接修改叶子节点B
new_B = B.clone()
# 替换目标位置，注意A需要扩展维度以匹配替换位置的数量
new_B[batch, replace_ids] = A.repeat(len(replace_ids), 1)

# 测试梯度传递
loss = new_B.sum()
loss.backward()
print(A.grad is not None)  # True
print(B.grad is not None)  # True

为什么之前的方法会失败 #

• 使用B[batch][replace_ids].data = A：直接修改.data属性会绕开PyTorch的自动求导跟踪机制，导致A的梯度无法传递到后续计算图，最终丢失梯度。
• 使用B[batch][replace_ids] = A：这是对叶子节点（B是requires_grad=True的叶子节点）的视图进行in-place赋值，PyTorch禁止此类操作——因为in-place修改会破坏计算图的完整性，导致反向传播时无法正确追踪梯度路径。

注意事项 #

• 如果原张量B是需要保留梯度的叶子节点，绝对不要直接修改原B，必须通过创建新张量（如clone或index_put生成的张量）参与后续计算。
• 替换时要保证A的维度与目标位置匹配：比如替换n个位置时，需要用A.repeat(n, 1)扩展A的维度，使其与B中被替换位置的维度一致。

内容的提问来源于stack exchange，提问作者Jometeorie