Skip to content

Commit

Permalink
[torchlib] Fix zeros_like signature (#1790)
Browse files Browse the repository at this point in the history
Include the missing arguments
  • Loading branch information
justinchuby authored Aug 7, 2024
1 parent cf5ddd9 commit b1f4942
Showing 1 changed file with 8 additions and 1 deletion.
9 changes: 8 additions & 1 deletion onnxscript/function_libs/torch_lib/ops/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -8960,7 +8960,14 @@ def aten_zeros(


@torch_op("aten::zeros_like", trace_only=True)
def aten_zeros_like(self: TTensor, dtype: int = -1) -> TTensor:
def aten_zeros_like(
self: TTensor,
dtype: int = -1,
layout: str = "",
device: str = "",
pin_memory: bool = False,
memory_format: str = "",
) -> TTensor:
"""zeros_like(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor"""

# NOTE: trace_only because both if branches need to be the same type, but we have
Expand Down

0 comments on commit b1f4942

Please sign in to comment.