PyTorch张量赋值异常：目标位置与实际修改位置不符求助

问题复现代码 #

import torch

# 创建CPU上的2×2零张量
empty_tensor = torch.zeros(2, 2)
# 创建MPS设备上的1D张量
r_tensor = torch.tensor([4.], device="mps")

# 执行赋值操作
empty_tensor[1, 1] = r_tensor
print(empty_tensor)
# 异常输出：
# tensor([[4., 0.],
#         [0., 0.]])

原因定位 #

这是PyTorch 2.1 nightly版本中MPS后端跨设备索引赋值的bug：当右侧赋值的张量位于MPS设备，而目标张量在CPU时，单元素索引（如[1,1]）的映射逻辑出现错误，导致赋值位置被错误映射到[0,0]。直接赋值Python float时，PyTorch会自动将其转换为CPU张量，避开了跨设备的索引解析问题，因此功能正常。

可行解决方案 #

• 方案1：将MPS张量转移到CPU后赋值

empty_tensor[1, 1] = r_tensor.cpu()

• 方案2：提取张量标量值赋值

empty_tensor[1, 1] = r_tensor.item()

• 方案3：统一张量设备
  将目标张量也移到MPS设备，避免跨设备操作：

empty_tensor = torch.zeros(2, 2, device="mps")
r_tensor = torch.tensor([4.], device="mps")
empty_tensor[1, 1] = r_tensor

补充说明 #

该问题属于nightly版本的后端实现bug，后续PyTorch正式版本大概率会修复。若需跟进，可在PyTorch GitHub仓库的issues中搜索相关问题提交反馈。

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