[export] Llama2 export with dynamic shapes fails using Torch 2.4 nightly about pytorch

[angelayi]

I wasn't able to repro this, but you could try doing this:

            import torch.export._trace

            ep = torch.export._trace._export(
                model, 
                (input_ids,), 
                dynamic_shapes=({1: seq_len},), 
                strict=False,
                _allow_complex_guards_as_runtime_asserts=True,
            )

The _allow_complex_guards_as_runtime_asserts flag should also allow you to export w/o needing torch.nn.attention.sdpa_kernel([SDPBackend.MATH])

[peri044]

I wasn't able to repro this

Is this using the reproducer provided or using the code snippet you suggested ?
Btw, the code snippet you provided seems to work fine but is it recommended to use torch.export._trace._export as a part of official examples ?