Repository: pytorch
Issue: Cannot export DETR model with torch.export
Labels: export-triaged,oncall: export

### 🐛 Describe the bug

Unable to export the DETR model encoder with torch.export

```python
# Fails
import torch
from torch.export import export
from transformers import DetrForObjectDetection

model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
encoder_inputs = (torch.rand(1, 850, 256), torch.rand(1, 850), torch.rand(1, 850, 256))
exported_mod = export(model.model.encoder, encoder_inputs)     ■ Cannot access member "
```

Gives the error:
```
lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1252, in guard_export_print
    assert (
AssertionError: whole graph export entails exactly one guard export
```

Please find in the following gist the corresponding `dynamo.explain` outputs for the same model (DETR encoder):
https://gist.github.com/corehalt/dbb031cd7403bb9fb9faa4bda025adf7


For reference, other parts of the model, as the backbone for example, succeeded:
```python
# Backbone is fine
import torch
from torch.export import export
from transformers import DetrForObjectDetection

model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
backbone_inputs = (torch.rand(1, 3, 800, 1066), torch.rand(1, 800, 1066))
exported_model = export(model.model.backbone, backbone_inputs)     ■ Cannot access membe
```

### Versions

Collecting environment information...                                                                                                                                                         
PyTorch version: 2.2.1+cpu                                                                                                                                                                    
Is debug build: False                                                                                                                                                                         
CUDA used to build PyTorch: None                                                                                                                                                              
ROCM used to build PyTorch: N/A                                                                                                                                                               
                                                                                                                                                                                              
OS: Ubuntu 22.04.3 LTS (x86_64)                                                                                                                                                               
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0                                                                                                                                            
Clang version: 17.0.6 (https://github.com/llvm/llvm-project.git 6009708b4367171ccdbf4b5905cb6a803753fe18)                                                                                     
CMake version: version 3.28.1                                                                                                                                                                 
Libc version: glibc-2.35                                                                                                                                                                      
                                                                                                                                                                                              
Python version: 3.10.4 (main, Mar  1 2024, 10:50:52) [GCC 11.4.0] (64-bit runtime)                                                                                                            
Python platform: Linux-5.15.0-91-generic-x86_64-with-glibc2.35                                                                                                                                
Is CUDA available: False                                                                                                                                                                      
CUDA runtime version: No CUDA                                                                                                                                                                 
CUDA_MODULE_LOADING set to: N/A                                                                                                                                                               
GPU models and configuration: No CUDA                                                                                                                                                         
Nvidia driver version: No CUDA                                                                                                                                                                
cuDNN version: No CUDA                                                                                                                                                                        
HIP runtime version: N/A                                                                                                                                                                      
MIOpen runtime version: N/A                                                                                                                                                                   
Is XNNPACK available: True                                                                                                                                                                    
                                                                                                                                                                                              
CPU:                                                                                                                                                                                          
Architecture:                       x86_64                                                                                                                                                    
CPU op-mode(s):                     32-bit, 64-bit                                                                                                                                            
Address sizes:                      52 bits physical, 57 bits virtual                                                                                                                         
Byte Order:                         Little Endian
CPU(s):                             256
On-line CPU(s) list:                0-255
Vendor ID:                          AuthenticAMD
Model name:                         AMD EPYC 9554 64-Core Processor
CPU family:                         25
Model:                              17
Thread(s) per core:                 2
Core(s) per socket:                 64
Socket(s):                          2
Stepping:                           1
Frequency boost:                    enabled
CPU max MHz:                        3762.9880
CPU min MHz:                        1500.0000
BogoMIPS:                           6200.13
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid overflow_recov succor smca fsrm flush_l1d
Virtualization:                     AMD-V
L1d cache:                          4 MiB (128 instances)
L1i cache:                          4 MiB (128 instances)
L2 cache:                           128 MiB (128 instances)
L3 cache:                           512 MiB (16 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-63,128-191
NUMA node1 CPU(s):                  64-127,192-255
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Mitigation; safe RET
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.3
[pip3] torch==2.2.1+cpu
[pip3] torchaudio==2.2.1+cpu
[pip3] torchvision==0.17.1+cpu
[conda] Could not collect

cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4