使用torch.cat跨设备指定out参数触发CUDA非法内存访问错误

• 运行环境：Jupyter Notebook，PyTorch 1.7.1
• 复现操作：调用torch.cat拼接两个CPU上的一维张量，同时指定一个位于CUDA设备的零值张量作为计算结果的输出载体

复现代码如下：

import torch
a = torch.ones(4)
b = torch.ones(4) 
c = torch.zeros(8).cuda()
print(c)
ab = torch.cat([a,b], out=c)
print(c)

运行后实际出现的报错栈：

...
\Anaconda3\envs\...\lib\site-packages\torch\_tensor_str.py in __init__(self, tensor)
     87 
     88         else:
---> 89             nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
     90 
     91             if nonzero_finite_vals.numel() == 0:

RuntimeError: CUDA error: an illegal memory access was encountered

• 版本校验逻辑遗漏：PyTorch 1.7.1版本的torch.cat在处理out参数时，没有对out张量和输入张量的设备一致性做前置显式检查，没有在函数调用入口直接抛出参数不匹配的可定位错误。
• 非法内存写入触发：当传入跨设备的out张量时，拼接算子会尝试将CPU端输入张量的计算结果，直接写入CUDA设备的显存地址空间。这类跨设备的内存写入没有合法的上下文校验，直接触发了CUDA运行时的非法内存访问错误。
• 报错位置错位的原因：CUDA的非法内存错误属于异步检测机制，不会在非法内存写入指令提交的当下立刻抛出异常，会延迟到后续下一次需要和CUDA上下文做同步交互的操作时才暴露。本例中torch.cat调用完成后没有触发CUDA同步，直到执行print(c)需要把CUDA张量的内容拉取到CPU端做字符串格式化时，才同步检测到之前的非法操作，因此报错栈的位置指向了打印逻辑，而非拼接逻辑本身。
• 文档描述的隐含要求：官方文档提及的“传入任意同类型张量组成的Python序列”，本身隐含了所有参与计算的张量（包括out参数指定的输出张量）必须满足同设备、同数据类型、形状匹配的一致性要求，只是早期版本没有将这类校验规则做明确的文档强调，也没有补全对应的前置检查逻辑。

• 所有带out参数的PyTorch算子调用，必须保证out张量的设备、数据类型、形状，和算子计算的预期输出完全匹配，禁止跨设备传入out参数。
• 针对本例的修正方案：要么提前将输入张量a、b迁移到CUDA设备再传入out=c，要么将输出张量c迁移到CPU设备后再作为out参数使用。

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