From b1f49425dd656a8f4cedbc527d3607867e3017b4 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 7 Aug 2024 16:32:32 -0700 Subject: [PATCH] [torchlib] Fix zeros_like signature (#1790) Include the missing arguments --- onnxscript/function_libs/torch_lib/ops/core.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 338ddb1ff..4b851b369 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -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