saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

attn_mask fix, set as torch.long (#18931)

Browse source 
### Description
<!-- Describe your changes. -->



### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
wejoncy 2023-12-26 16:43:33 +08:00 committed by GitHub
parent 37f743680a
commit 0c5826f013
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
onnxruntime/python/tools/transformers/large_model_exporter.py
View file

@ -252,7 +252,8 @@ def fetch_onnx_inputs_outputs_name(
elif name == "attention_mask":
attn_mask = onnx_inputs[idx]
onnx_inputs[idx] = torch.cat(
(attn_mask, torch.ones((attn_mask.shape[0], 1), device=attn_mask.device)), dim=1
(attn_mask, torch.ones((attn_mask.shape[0], 1), device=attn_mask.device, dtype=attn_mask.dtype)),
dim=1,
)
elif name == "input_ids":
input_ids = onnx_inputs[idx]