From 8a454bc695f9f0e035ebc3f7fc23f4a2e91f4651 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 20 Jan 2025 07:18:36 -0800 Subject: [PATCH] [torchlib] Use Optional in type annotations for linear --- onnxscript/function_libs/torch_lib/ops/nn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index d91a12ec35..8bb8bf0aa3 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -826,7 +826,7 @@ def aten_leaky_relu_backward( @torch_op("aten::linear", trace_only=True) -def aten_linear(input: TFloat, weight: TFloat, bias: TFloat | None = None) -> TFloat: +def aten_linear(input: TFloat, weight: TFloat, bias: Optional[TFloat] = None) -> TFloat: """linear(Tensor input, Tensor weight, Tensor? bias=None) -> Tensor""" if len(input.shape) == 2: