Skip to content

Commit

Permalink
add TORCHDYNAMO_CAPTURE_DYNAMIC_OUTPUT_SHAPE_OPS (#127017)
Browse files Browse the repository at this point in the history
tlparse prints failure description like this

> dynamic shape operator: aten._unique2.default; to enable, set torch._dynamo.config.capture_dynamic_output_shape_ops = True

adding os env var to set it easier for testing

Pull Request resolved: #127017
Approved by: https://github.com/jackiexu1992
  • Loading branch information
dshi7 authored and pytorchmergebot committed May 25, 2024
1 parent 84e59f0 commit 0f67d38
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion torch/_dynamo/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -221,7 +221,9 @@ def is_fbcode():
# break instead of capturing. This requires dynamic_shapes to be True.
# If you set this to True, you probably also want capture_scalar_outputs
# (these are separated for historical reasons).
capture_dynamic_output_shape_ops = False
capture_dynamic_output_shape_ops = (
os.environ.get("TORCHDYNAMO_CAPTURE_DYNAMIC_OUTPUT_SHAPE_OPS", "0") == "1"
)

# By default, dynamo will treat all ints as backed SymInts, which means (1) it
# will wait to see the int change over multiple runs before generalizing and
Expand Down

0 comments on commit 0f67d38

Please sign in to comment.