pytorch - 💡(How to fix) Fix [ONNX] ExportedProgram with register_buffer + in-place copy fails on torch.onnx.export [1 comments, 2 participants]

ValueError: Key 'b_prompt_feat' does not match the name of the value 'copy'.
Please use the value.name as the key.

---

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer('prompt_feat', torch.zeros(1, 4, 8))
        self.linear = nn.Linear(8, 8)

    def forward(self, x):
        out = torch.zeros(1, 4, 8)
        out[:, :, :] = self.prompt_feat.to(x.device)
        return self.linear(out + x)

model = Model().eval()
x = torch.randn(1, 4, 8)

ep = torch.export.export(model, (x,))
torch.onnx.export(ep, args=(x,), f="test.onnx")  # ValueError

---

graph_signature.inputs_to_buffers: 'b_prompt_feat' → 'prompt_feat'

---

[b_prompt_feat]  placeholder [1, 4, 8] float32
[slice]          aten.slice(zeros, ...)
[copy]           aten.copy.default(slice, b_prompt_feat)   ← both float32
[slice_scatter]  aten.slice_scatter(zeros, copy, ...)

---

# torchlib: aten_copy
@torch_op("aten::copy", trace_only=True)
def aten_copy(self, src, non_blocking=False):
    return op.CastLike(src, self)

---

# line ~139-141
node_name_to_values[node.name] = outputs   # values['copy'] = value_A
outputs.name = node.name                    # value_A.name = 'copy'  ← BUG

---

# _exported_program_to_onnx_program, line ~140
model.graph.initializers['b_prompt_feat'] = value
# key='b_prompt_feat' but value.name='copy' → ValueError

---

# Before (line ~141):
outputs.name = node.name

# After:
if outputs.producer() is not None:
    # Output is a new value produced by a new ONNX node — safe to rename
    outputs.name = node.name
# else: output is an existing value (identity passthrough) — do not rename