Skip to content

Commit

Permalink
Fix missing type in _add_attribute_to_torchscript_node for Deberta mo…
Browse files Browse the repository at this point in the history
…dels (#1773)

Signed-off-by: Xavier Dupre <[email protected]>
  • Loading branch information
xadupre authored Aug 7, 2024
1 parent 66e6a66 commit 2b5173d
Showing 1 changed file with 12 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -459,8 +459,18 @@ def _add_attribute_to_torchscript_node(
return node.fs_(key, list(value)) # type: ignore[arg-type]
if isinstance(value[0], int):
return node.is_(key, list(value)) # type: ignore[attr-defined]
raise TypeError(f"Unsupported sequence type '{type(value)}' for attribute '{key}'")
raise TypeError(f"Unsupported attribute type '{type(value)}' for attribute '{key}'")
raise TypeError(
f"Unsupported sequence type '{type(value)}' for attribute '{key}' in "
f"node={node!r}, value is {value!r}"
)
if "TensorProtoDataType" in str(type(value)):
# torch._C._onnx.TensorProtoDataType
return node.i_(key, int(value))

raise TypeError(
f"Unsupported attribute type '{type(value)}' for attribute '{key}' "
f"in node={node!r}, value is {value!r}"
)


@runtime_typing.checked
Expand Down

0 comments on commit 2b5173d

Please sign in to comment.