python export-onnx.py small PyTorch 2.0+

[kendiyang]

Traceback (most recent call last):
File "/Users/kendi/temp/sherpa-onnx/scripts/whisper/export-onnx.py", line 609, in
main()
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/temp/sherpa-onnx/scripts/whisper/export-onnx.py", line 545, in main
torch.onnx.export(
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/onnx/init.py", line 375, in export
export(
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/onnx/utils.py", line 502, in export
_export(
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/onnx/utils.py", line 1564, in _export
graph, params_dict, torch_out = _model_to_graph(
^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/onnx/utils.py", line 1113, in _model_to_graph
graph, params, torch_out, module = _create_jit_graph(model, args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/onnx/utils.py", line 997, in _create_jit_graph
graph, torch_out = _trace_and_get_graph_from_model(model, args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/onnx/utils.py", line 904, in _trace_and_get_graph_from_model
trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/jit/_trace.py", line 1500, in _get_trace_graph
outs = ONNXTracedModule(
^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/jit/_trace.py", line 139, in forward
graph, out = torch._C._create_graph_by_tracing(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/jit/_trace.py", line 130, in wrapper
outs.append(self.inner(*trace_inputs))
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1726, in _slow_forward
result = self.forward(*input, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/temp/sherpa-onnx/scripts/whisper/export-onnx.py", line 249, in forward
x, self_k_cache, self_v_cache = block(
^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1726, in _slow_forward
result = self.forward(*input, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/temp/sherpa-onnx/scripts/whisper/export-onnx.py", line 204, in forward
self_attn_x, self_k_cache_updated, self_v_cache_updated = self.attn(
^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1726, in _slow_forward
result = self.forward(*input, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/temp/sherpa-onnx/scripts/whisper/export-onnx.py", line 180, in forward
wv, qk = self.multiHeadAttention.qkv_attention(q, k_cache, v_cache, mask)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kendi/.venv/lib/python3.12/site-packages/whisper/model.py", line 124, in qkv_attention
a = scaled_dot_product_attention(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: scaled_dot_product_attention(): argument 'is_causal' must be bool, not Tensor

[csukuangfj]

PyTorch 2.0+

Could you be more specific?

Do you mean 2.0.0, or 2.0.1, or 2.2.0, or 2.5.1?