void onnxruntime::BroadcastIterator::Init(ptrdiff_t, ptrdiff_t) axis == 1 || axis == largest was false. Attempting to broadcast an axis by a dimension other than 1. 3 by 224

[MHGL]

🐛 Bug

I get an error

• reigster_buffer in Module
• onnxruntime on onnx from torch.onnx.export

To Reproduce

Steps to reproduce the behavior:

1. code example

import torch

# init module
class MyModule(torch.nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.register_buffer("rb", torch.randn(1, 1, 3, 1, 1)) 
        # you can run with the workaround
        # self.rb = torch.randn(1, 1, 3, 1, 1)

    def forward(self, x):
        x += self.rb[0]
        return x

torch_model = MyModule().eval()

# torch.onnx.export
torch.onnx.export(torch_model,
        torch.randn(1, 3, 224, 224),
        "./tmp.onnx",
        input_names=["inputs"],
        output_names=["outputs"],
        dynamic_axes={"inputs": {0: "batch", 2: "height", 3: "width"}, "outputs": {0: "batch", 1: "class", 2: "height", 3: "width"}},
        opset_version=11,
        export_params=True)

# onnxruntime
import os
import numpy as np
import onnxruntime
from onnxruntime.datasets import get_example

onnx_model = get_example(os.path.join(os.getcwd(), "tmp.onnx"))
sess = onnxruntime.InferenceSession(onnx_model)
inputs = np.random.randn(1, 3, 224, 224).astype(np.float32)
onnx_out = sess.run(None, {"inputs": inputs})

2. stack traces

Warning: ONNX Preprocess - Removing mutation from node aten::add_ on block input: '0'. This changes graph semantics.
2021-07-07 20:06:52.126767616 [E:onnxruntime:, sequential_executor.cc:339 Execute] Non-zero status code returned while running Add node. Name:'Add_0' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/math/element_wise_ops.h:497 void onnxruntime::BroadcastIterator::Init(ptrdiff_t, ptrdiff_t) axis == 1 || axis == largest was false. Attempting to broadcast an axis by a dimension other than 1. 3 by 224

Traceback (most recent call last):
  File "mini_code.py", line 37, in <module>
    onnx_out = sess.run(None, {"inputs": inputs})
  File "/opt/conda/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 188, in run
    return self._sess.run(output_names, input_feed, run_options)
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Add node. Name:'Add_0' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/math/element_wise_ops.h:497 void onnxruntime::BroadcastIterator::Init(ptrdiff_t, ptrdiff_t) axis == 1 || axis == largest was false. Attempting to broadcast an axis by a dimension other than 1. 3 by 224

Expected behavior

Environment

• PyTorch Version: 1.9.0
• OS (e.g., MacOS, Linux): Ubuntu20.04 LTS
• How you install python (anaconda, virtualenv, system): miniconda
• python version (e.g. 3.7): 3.8.5
• any other relevant information:
  • gpu: GeForce GTX 1650
  • driver: Driver Version: 460.80
  • CUDA: CUDA Version: 11.2

[lyhue1991]

@MHGL have you solved this problem?