先说结论，这个报错不是因为你代码写得有问题，而是因为你在把 PyTorch 模型导出成 ONNX 时，模型里某个操作触发了aten::index_put_算子，而 PyTorch 的 ONNX 导出器对这个算子的支持非常有限。我前阵子做一个需要部署到服务端的文本分类模型时，也卡在这个报错上整整一个晚上。
报错原文大概是这样的：
RuntimeError: Only consecutive 1-d tensor indices are supported in exporting aten::index_put to ONNX如果你也遇到了，千万别急着重写模型。这篇文章我就用自己实际踩坑的经历，把这个报错从根因到替代方案完整梳理一遍，最后还会给你一份排查清单，照着做基本能解决大部分类似问题。
1. 先搞明白报错在说什么：ONNX 导出过程的"花式拒绝"
我最早看到这个报错时，第一反应是去搜aten::index_put是什么意思。这个算子的全称是 ATen 里的index put操作，底层对应的就是你写x[mask] = value或者x[indices] = value这类"按索引赋值"的代码。
举个最简单的例子：
x = torch.zeros(5) mask = torch.tensor([True, False, True, False, True]) x[mask] = 1.0这两行代码在 PyTorch 内部就会触发index_put_（带下划线表示 in-place 操作）。平时训练、推理都没问题，但在导出 ONNX 时，导出器需要把 PyTorch 算子逐个映射到 ONNX 算子。问题就出在这个映射上。
ONNX 里其实是有ScatterElements和ScatterND这类算子的，理论上可以表达"按索引赋值"。但 PyTorch ONNX 导出器对index_put_的支持只覆盖了一个非常狭窄的场景：连续的 1 维索引。
什么叫连续 1 维索引？就是x[2:5] = value这种切片式赋值，索引是一段连续区间。报错里提到的 "consecutive 1-d tensor indices"，指的就是x[start:end]或者x[[start, start+1, ..., end-1]]这类可以用连续区间表示的下标。
而你一旦写了x[mask] = value这种布尔掩码赋值，或者x[[1, 5, 7]] = value这种任意位置组合的赋值，导出器就直接撂挑子不干了。
1.1 为什么模型里会"莫名其妙"出现这个算子
这是最坑的地方。大部分情况下，你不会直接在模型前向代码里写x[mask] = value，但它会通过其他高层 API 间接触发。
我那次就是在一个自注意力模块里用了logits.masked_fill_(mask, -1e9)，本意是把注意力分数里某些位置填充成一个极小的数。你猜怎么着？masked_fill_在底层也会走index_put_的路径。
类似的 API 包括但不限于：
Tensor.masked_fill()/masked_fill_()Tensor.index_fill()/index_fill_()Tensor.scatter()/scatter_()- 自定义损失函数里的
output[target] = 0这类写法 - 某些第三方库的 mask 操作，比如
transformers里部分填充逻辑
所以你会发现一个问题：代码里明明没写任何"危险操作"，导出时照样报这个错。原因就是高层 API 隐藏了底层算子的真实面貌，报错时你根本不知道是哪一行代码闯的祸。
1.2 这个报错的时间线：为什么"以前没事，现在出事"
还有一个现象值得注意：同一个模型，不同 PyTorch 版本下导出结果可能完全不同。
我试过的版本差异大致是这样：
| PyTorch 版本 | ONNX 导出行为 |
|---|---|
| 1.8 及更早 | 部分场景直接报错或导出成不正确的图 |
| 1.9 ~ 1.12 | 严格检查，遇到非连续索引直接报 RuntimeError |
| 2.0 以上（TorchScript 路径） | 行为不稳定，有的能导出但推理结果错误 |
| 2.1+（dynamo 路径） | 默认导出器仍受限，但暴露了更多内部细节 |
这个报错在 1.x 时代就存在，到现在也没彻底消失，根本原因是 ONNX 的算子集和 PyTorch 的算子集本来就不是一一对应的。PyTorch 的算子粒度很细，index_put_这种"写入"+ "索引组合" 的语义在 ONNX 里没有天然对应的单一算子，只能靠导出器做模式匹配。匹配不上就报错，简单粗暴。
所以你真正要做的，不是研究怎么让导出器支持index_put_，而是把模型里的索引赋值操作改写成 ONNX 能理解的形式。
2. 定位肇事代码：三步锁定"元凶"行
光知道原理还不够，你得知道具体是哪一行代码在导出时炸的。我那次排查花了大半晚上，最后总结出了一套定位流程，效率高很多，这里分享给你。
2.1 第一步：用最小复现脚本确认触发条件
不要一上来就在完整模型上排查，先写一个最小脚本，确认报错能否稳定复现，同时验证你的修改方向。比如：
import torch import torch.onnx class IndexPutModel(torch.nn.Module): def forward(self, x, mask): # 触发 index_put_ 的操作 x = x.clone() x[mask] = 0.0 return x model = IndexPutModel() dummy_input = torch.randn(8, 10) dummy_mask = torch.tensor([True, False, True, False, True, False, True, False]) torch.onnx.export( model, (dummy_input, dummy_mask), "test.onnx", opset_version=13, input_names=["x", "mask"], )在我印象里，这个脚本稍微改一改，比如把x[mask] = 0.0换成x[0:3] = 0.0，就能正常导出。用这个最小脚本反复验证，比在完整模型上瞎猜高效得多。
2.2 第二步：用二分注释法缩小范围
如果最小复现脚本没报错，说明问题出在模型内部的某个模块里。这时我惯用的方法非常土但非常有效：先把模型整体的导出脚本写好，然后注释掉模型代码的一半，导出试试；如果还报错，再注释掉剩下的一半的一半。
具体操作：
- 把
model.forward()里每个大模块单独导出一次，找到报错的那个模块。 - 进入报错模块，逐个子模块嵌套导出，继续二分。
- 找到具体模块后，再检查模块内的数据流，搜索
masked_fill、scatter、index_fill、[...] = ...这类关键字。
最后锁定可疑代码后，把输入固定好，单独跑一遍导出，确认报错确实由它引起。
这个方法听着笨，但每次都很灵验。尤其是模型结构复杂、嵌套了好几层nn.Module的时候，靠肉眼直接看代码经常会漏掉藏在某个子模块里的暗桩。
2.3 第三步：利用报错堆栈定位
有时候报错信息会附带 Python 调用栈，虽然不精确，但能指出方向。
比如堆栈里出现：
File "torch/onnx/utils.py", line XXXX, in _export graph, _, _ = model(*args) File "torch/onnx/symbolic_opset16.py", line XXX, in index_put如果你的代码没有任何封装、纯nn.Module直接导出，那堆栈里会出现model.forward的调用位置，顺着找就行。但注意，如果模型里用了torch.jit.trace或torch.jit.script，堆栈可能被消化掉，只显示torch/onnx内部的帧，这时就必须回到二分法。
补充一个我后来养成的习惯：写一个公共导出脚本，支持传入模块名字符串，自动尝试导出指定子模块，这样定位下一轮问题时不用反复改代码。
3. ONNX 友好改造的核心方案：把索引赋值改成 torch.where 组合
定位到肇事代码之后，接下来的问题就是怎么改。先说结论：在绝大多数场景下，用torch.where替代索引赋值，是改动最小、语义最接近、导出成功率最高的方案。
3.1 torch.where 的替代逻辑
假设你原来的代码是：
x[mask] = value它表达的意思是：把x中mask为真的位置替换成value，其余位置保留原值。这本质上就是一个逐元素的条件选择。而torch.where干的正是这件事：
x = torch.where(mask, value, x)torch.where(condition, a, b)对每个位置做判断：condition为真取a，为假取b。当a是一个标量时，PyTorch 会自动广播到和x相同的形状。ONNX 导出器对torch.where的支持非常成熟，因为它直接映射到 ONNX 的Where算子。
以我那个masked_fill_(mask, -1e9)为例，改写前是：
logits.masked_fill_(mask, -1e9)改写后是：
logits = torch.where(mask, torch.full_like(logits, -1e9), logits)如果mask的 shape 和logits不一致，torch.where的广播机制也能处理，只要它们能广播到同一形状。这一点比直接索引赋值还灵活。
3.2 几种典型场景的具体改写模板
我根据自己实际遇到的情况，整理了几个高频场景的改写模板。
场景一：标量替换
原代码：
x[mask] = 0.0改写：
x = torch.where(mask, torch.zeros_like(x), x)注意zeros_like保持了和x完全相同的 dtype 和 device，避免类型不匹配。
场景二：用另一个张量的值替换
原代码：
x[mask] = y[mask]改写：
x = torch.where(mask, y, x)只有mask为真的位置才取y的同位置值，其余保留x原值。如果y和x形状不同但可广播，也能正常工作。
场景三：二维索引赋值
原代码：
x[:, idx] = scalar_value这在 NLP 模型里很常见，比如把某个 token 位置的 logit 强制设为极大或极小值。改写时可以先用torch.arange构造一个 mask：
row_indices = torch.arange(x.shape[1], device=x.device) mask = (row_indices == idx) # shape: [seq_len] x = torch.where(mask.unsqueeze(0), torch.full_like(x, scalar_value), x)如果要对 batch 内每个样本的不同位置赋值，那就把 mask 构造成[batch_size, seq_len]的形状，torch.where一样能处理。
3.3 改写后的边界检查
改写不是替换完就完事，有几个边界问题你必须检查：
第一个：dtype 一致性。torch.full_like(x, value)里的value会按x的 dtype 解释。如果x是 int64、而你想填的value是负数-1e9，那结果会溢出或直接截断。所以填充大负数时先确认x是浮点类型。
第二个：in-place 语义变化。原来的x[mask] = value是 in-place 操作，x本身被修改；改写后x = torch.where(...)会创建一个新张量。如果你的代码是logits.masked_fill_(mask, -1e9)这种带下划线的 in-place 版本，改写时记得把后续所有对logits的引用都理解成"指向新张量"，不要在前面保留旧引用，否则后续计算用的还是旧值。
第三个：mask 的 dtype。torch.where的condition必须是 bool 张量。如果你的 mask 是 0/1 的 float 张量（有些人用(x > 0).float()来生成 mask），需要先转成 bool：
mask_bool = mask.to(torch.bool) # 或者 mask > 0这个细节很容易被忽略，踩过一次你就记住了。
4. 其他可用的替代方案：masked_fill、scatter、自定义算子
torch.where不是唯一的路。根据具体场景，下面几种方案也值得掌握。
4.1 masked_fill 在某些场景下依然可用
你可能会问：既然masked_fill_触发index_put_，那masked_fill（不带下划线）呢？
我实测下来的结论是：部分场景可以导出，但稳定性没有保证。在特定 opset 版本、特定 mask 形状下，masked_fill的导出路径可能走的是Where而不是index_put_，这时没问题。但换个 mask 形状，可能又打回原形，报index_put_的错误。
所以我的建议是：如果你已经在排查这个报错了，就直接用torch.where，别在masked_fill和masked_fill_之间反复横跳浪费时间。你的目标是成功导出，不是帮 PyTorch 测试导出器的算子覆盖范围。
4.2 scatter / index_fill 系列：能规避就规避
scatter、index_fill、index_copy这几个 API 底层同样会进入index_put_的领地，ONNX 导出支持也分场景：
scatter在 opset 11+ 里有ScatterElements对应，但语义不完全一致，容易出岔子。index_fill的导出支持也很窄，只处理连续索引等简单情形。index_copy在多数版本里导出直接失败。
如果你真的需要按索引把某个位置替换成固定值，还是建议回到torch.where方案。
4.3 自定义算子：最后的手段
有些场景确实绕不开index_put_的语义，比如你要在 ONNX 图里对一个动态位置做复杂写入。这时就得上自定义算子。
完整流程如下：
- 在 PyTorch 里用
torch.library注册一个自定义算子。 - 编写该算子的 symbolic function，告诉 ONNX 导出器映射到哪个 ONNX 算子（比如
ScatterND）。 - 在推理引擎（如 ONNX Runtime）里注册对应的 custom op 实现。
这个方案工程量不小，还要在部署端配套开发，不到万不得已我不推荐。如果你只是导出模型给 TensorRT、ONNX Runtime 这类常见引擎用，torch.where几乎总能帮你绕过去。
4.4 新版导出路径 dynamo=True 的差异
PyTorch 2.x 提供了基于 dynamo 的新导出路径，写法是：
torch.onnx.export(model, args, "model.onnx", dynamo=True)这个新路径理论上对更多算子有更完整的支持，但我实测下来，index_put_的限制仍然以不同形式存在。有时候报错信息会从RuntimeError变成别的，比如：
Unsupported: ONNX export of operator index_put本质没变，只是换了层皮。所以在导出用于生产环境的 ONNX 模型时，我还是以"修改模型代码让它对 ONNX 友好"为主，不指望导出器突然变得万能。
5. 避坑清单：ONNX 导出前先自查这几类操作
经历了那一个晚上的折腾，我后来养成了一个习惯：在写任何要导出 ONNX 的模型前，先自查代码里是否有下面这几类"雷区操作"。提前规避比事后排查省太多时间。
5.1 高频雷区操作清单
| 操作类型 | 示例代码 | ONNX 导出风险 |
|---|---|---|
| 布尔掩码赋值 | x[mask] = value | 极高 |
| 任意索引赋值 | x[[1,3,5]] = value | 极高 |
| 二维行列索引赋值 | x[:, idx] = value | 极高 |
| masked_fill_ | x.masked_fill_(mask, value) | 高，取决于 mask 形状 |
| scatter_ | x.scatter_(dim, index, src) | 高，opset 版本敏感 |
| index_fill_ | x.index_fill_(dim, index, value) | 高 |
| in-place 索引更新 | x[condition] += 1 | 极高 |
如果你写完代码发现上面这些一个都没有，恭喜你，踩雷概率大幅降低。
5.2 我常用的"导出前扫描脚本"
为了不每次都手动检查，我写过一个简单的 AST 扫描脚本，代码量不大，但很实用。原理是遍历 Python 源码的 AST，找出所有Subscript+Assign的组合，以及masked_fill、scatter、index_fill等函数调用。
精简版思路如下：
import ast class IndexPutVisitor(ast.NodeVisitor): def visit_Assign(self, node): # 检查 x[mask] = value 这类赋值 if isinstance(node.targets[0], ast.Subscript): print(f"可疑索引赋值: {ast.dump(node)}") self.generic_visit(node) def visit_Call(self, node): # 检查 masked_fill / scatter / index_fill 调用 if isinstance(node.func, ast.Attribute): name = node.func.attr if name in {"masked_fill", "masked_fill_", "scatter", "scatter_", "index_fill", "index_fill_", "index_put_"}: print(f"可疑API调用: {name}") self.generic_visit(node) with open("model.py") as f: tree = ast.parse(f.read()) IndexPutVisitor().visit(tree)这个脚本的作用不是拦截所有问题，而是帮你快速圈定可疑区域。排查时直接跳转到这些行附近检查，不用满文件找。
5.3 导出前数值验证：避免"导出成功但结果不对"
最后再提醒一个容易忽略的点：不要以为不报错就等于万事大吉。
ONNX 导出有成功但语义漂移的可能。index_put_这个场景尤其容易踩坑，因为有些版本虽然导出成功，但生成的 ONNX 图里计算逻辑已经和 PyTorch 不一致了。
所以我每次改写代码后，都会做一遍数值验证：
import numpy as np import onnxruntime as ort # PyTorch 推理结果 torch_output = model(dummy_input).detach().numpy() # ONNX Runtime 推理结果 sess = ort.InferenceSession("test.onnx") onnx_output = sess.run(None, {"x": dummy_input.numpy(), "mask": dummy_mask.numpy()})[0] # 对比 print("最大绝对误差:", np.abs(torch_output - onnx_output).max()) print("是否一致:", np.allclose(torch_output, onnx_output, atol=1e-6))如果你的操作涉及负数填充、极值填充，atol可以放宽一些，但绝对误差不超过1e-4。如果误差超了，说明改写后的语义和原来有出入，需要回头检查广播和 dtype 问题。
我第一次改写注意力 mask 时，就是因为mask的 shape 和logits不一致但又能隐式广播，导致 ONNX 里的Where和 PyTorch 里的index_put_语义产生了微妙差异。那一次要不是跑了数值验证，模型上线后推理结果可能悄悄漂移，后果不堪设想。
5.4 最后的小技巧：不要遗忘版本记录
排查这类问题过程中，我越来越觉得"记录环境版本 + 操作代码"很重要。ONNX 导出的行为与 PyTorch 版本、opset 版本、甚至操作系统都强相关。我建议你在项目里维护一个小文件，记录：
- PyTorch 版本
- opset_version
- 导出时用的
torch.onnx.export参数 - 报错信息原文
- 最终采用的改写方案
以后项目升级或换机器重新部署时，这份记录能帮你省下从头排查的时间。这个习惯是我在连续踩了几个 ONNX 相关坑之后才养成的，现在想想，早该如此。
回到开头那个报错。其实它并不难解决，难的是搞清楚两件事：底层的算子语义是什么，以及导出器能接受的替代写法是什么。搞明白了这两点，你会发现RuntimeError: Only consecutive 1-d tensor indices...只是一个善意的提醒——它告诉你，该让模型代码更贴近 ONNX 的逻辑了。而torch.where这条改造路径，到目前为止，是我用过的最顺手、最稳妥的一招。
