From 377869a7720505a78985520ad6e735ac53c7d0b6 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 13 Sep 2024 08:58:36 -0700 Subject: [PATCH] [torch api] Enable `_TORCH_ONNX_SAVE_EXTERNAL_DATA_WITH_IR` (#1866) Enable the `_TORCH_ONNX_SAVE_EXTERNAL_DATA_WITH_IR` flag to use the new external data logic. This will - Reduce peak memory usage - Align external data to 64k for the torch exporter. --- onnxscript/_framework_apis/torch_2_5.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxscript/_framework_apis/torch_2_5.py b/onnxscript/_framework_apis/torch_2_5.py index 6d458bc65..d011e0a17 100644 --- a/onnxscript/_framework_apis/torch_2_5.py +++ b/onnxscript/_framework_apis/torch_2_5.py @@ -25,7 +25,7 @@ # Internal flag. Will go away. _TORCH_ONNX_SAVE_EXTERNAL_DATA_WITH_IR = ( - os.getenv("TORCH_ONNX_OFFLOAD_EXTERNAL_DATA_WITH_IR") == "1" + os.getenv("TORCH_ONNX_OFFLOAD_EXTERNAL_DATA_WITH_IR") != "0" )