PyTorch Transformer开发触发张量跨设备RuntimeError问题

逐位与运算的两个输入张量trg_mask、np_mask分属不同计算设备（分别位于cuda:0和CPU上），PyTorch不支持跨设备执行逐元素运算，因此触发设备不匹配异常。

• 设备判断逻辑完全依赖手动传入的opt.device硬编码值，没有和实际输入张量的设备状态做动态同步：只要配置项和运行时的张量位置不一致（比如手动把输入张量搬到CUDA但没更新opt参数），生成的nopeak掩码就会留在CPU上，和已经在CUDA上的trg_mask形成设备错位
• 原代码中被注释掉的CUDA搬运逻辑存在API使用错误：PyTorch的张量.cuda()方法是返回迁移后的新张量，不会原地修改原变量，就算取消注释也无法完成设备搬运
• 缺少对掩码张量的设备对齐机制，只要运行环境和预设配置有偏差就会触发报错

import torch
import numpy as np

def nopeak_mask(size, ref_tensor):
    # 生成因果掩码基础结构
    np_mask = np.triu(np.ones((1, size, size)), k=1).astype('uint8')
    np_mask = torch.from_numpy(np_mask) == 0
    # 直接跟随参考张量的所在设备生成掩码，避免手动判断偏差
    return np_mask.to(ref_tensor.device)

def create_masks(src, trg, opt):
    src_mask = (src != opt.src_pad).unsqueeze(-2)
    if trg is not None:
        trg_mask = (trg != opt.trg_pad).unsqueeze(-2)
        size = trg.size(1)
        # 以trg张量的设备为基准生成nopeak掩码，保证两个参与运算的张量设备完全一致
        np_mask = nopeak_mask(size, trg)
        trg_mask = trg_mask & np_mask
    else:
        trg_mask = None
    return src_mask, trg_mask

• 完全摆脱对固定配置参数的设备判断依赖，不管输入张量运行在CPU还是CUDA设备上，生成的掩码都会自动对齐到对应设备
• 替换了原代码中无效的CUDA搬运逻辑，不存在API使用错误
• 从机制层面杜绝跨设备运算的触发可能

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