From 1efd8e665159bb3fe170628a923db435e165942f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 3 Jun 2024 16:36:16 -0700 Subject: [PATCH] Remove traceable from `_aten_scaled_dot_product_efficient_attention_fillin_empty_outputs` (#1585) Fix _aten_scaled_dot_product_efficient_attention_fillin_empty_outputs regression Co-authored-by: Ti-Tai Wang --- 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 7730008ef..ea16f4c37 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -1779,7 +1779,7 @@ def aten__scaled_dot_product_flash_attention( ) -@torch_op("aten::_scaled_dot_product_efficient_attention", private=True, traceable=True) +@torch_op("aten::_scaled_dot_product_efficient_attention", private=True) def _aten_scaled_dot_product_efficient_attention_fillin_empty_outputs( query: TFloat, compute_log_sumexp: bool,