Repository: pytorch
Issue: Error while exporting `keypointrcnn_resnet50_fpn`
Labels: oncall: pt2,oncall: export

### 🐛 Describe the bug



Code for reproducing the problem

```
from torchvision.models.detection import keypointrcnn_resnet50_fpn, KeypointRCNN_ResNet50_FPN_Weights
import torch

model = keypointrcnn_resnet50_fpn(weights=KeypointRCNN_ResNet50_FPN_Weights.DEFAULT)
model.eval()

input_frame = torch.randn(1,3, 4898, 3265)

exported_program: torch.export.ExportedProgram= torch.export.export(model, args=(input_frame,), strict=False)

```

Error logs

```
W1015 22:00:17.836165 3292 site-packages/torch/fx/experimental/symbolic_shapes.py:6061] failed during evaluate_expr(4*u1 > 4000, hint=None, size_oblivious=False, forcing_spec=False
E1015 22:00:17.836511 3292 site-packages/torch/fx/experimental/recording.py:298] failed while running evaluate_expr(*(4*u1 > 4000, None), **{'fx_node': False})
Traceback (most recent call last):
  File "/home/ubuntu/experiments/export/keyrcnn_pose.py", line 9, in <module>
    exported_program: torch.export.ExportedProgram= torch.export.export(model, args=(input_frame,), strict=False)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/export/__init__.py", line 366, in export
    return _export(
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/export/_trace.py", line 1014, in wrapper
    raise e
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/export/_trace.py", line 987, in wrapper
    ep = fn(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/export/exported_program.py", line 116, in wrapper
    return fn(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/export/_trace.py", line 1964, in _export
    export_artifact = export_func(  # type: ignore[operator]
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/export/_trace.py", line 1754, in _non_strict_export
    aten_export_artifact = _to_aten_func(  # type: ignore[operator]
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/export/_trace.py", line 643, in _export_to_aten_ir
    gm, graph_signature = transform(aot_export_module)(
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/export/_trace.py", line 1684, in _aot_export_non_strict
    gm, sig = aot_export(wrapped_mod, args, kwargs=kwargs, **flags)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 1262, in aot_export_module
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 1497, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 524, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 625, in _create_aot_dispatcher_function
    fw_metadata = run_functionalized_fw_and_collect_metadata(
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/collect_metadata_analysis.py", line 194, in inner
    flat_f_outs = f(*flat_f_args)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 184, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 863, in functional_call
    out = mod(*args[params_len:], **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/export/_trace.py", line 1671, in forward
    tree_out = mod(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torchvision/models/detection/generalized_rcnn.py", line 104, in forward
    proposals, proposal_losses = self.rpn(images, features, targets)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torchvision/models/detection/rpn.py", line 373, in forward
    boxes, scores = self.filter_proposals(proposals, objectness, images.image_sizes, num_anchors_per_level)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torchvision/models/detection/rpn.py", line 289, in filter_proposals
    keep = box_ops.batched_nms(boxes, scores, lvl, self.nms_thresh)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torchvision/ops/boxes.py", line 72, in batched_nms
    if boxes.numel() > (4000 if boxes.device.type == "cpu" else 20000) and not torchvision._is_tracing():
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/__init__.py", line 686, in __bool__
    return self.node.bool_()
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/fx/experimental/sym_node.py", line 511, in bool_
    return self.guard_bool("", 0)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/fx/experimental/sym_node.py", line 449, in guard_bool
    r = self.shape_env.evaluate_expr(self.expr, self.hint, fx_node=self.fx_node)
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/fx/experimental/recording.py", line 262, in wrapper
    return retlog(fn(*args, **kwargs))
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 6057, in evaluate_expr
    return self._evaluate_expr(
  File "/home/ubuntu/anaconda3/envs/export/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 6207, in _evaluate_expr
    raise self._make_data_dependent_error(
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not guard on data-dependent expression 4*u1 > 4000 (unhinted: 4*u1 > 4000).  (Size-like symbols: u1)

Caused by: (torchvision/ops/boxes.py:72 in batched_nms)
```

### Versions

```
PyTorch version: 2.6.0.dev20241015+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.10.0 | packaged by conda-forge | (default, Nov 20 2021, 02:24:10) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-1063-aws-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: GPU 0: NVIDIA A10G
Nvidia driver version: 535.183.01
cuDNN version: Could not collect
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
Byte Order:                         Little Endian
Address sizes:                      48 bits physical, 48 bits virtual
CPU(s):                             8
On-line CPU(s) list:                0-7
Thread(s) per core:                 2
Core(s) per socket:                 4
Socket(s):                          1
NUMA node(s):                       1
Vendor ID:                          AuthenticAMD
CPU family:                         23
Model:                              49
Model name:                         AMD EPYC 7R32
Stepping:                           0
CPU MHz:                            2800.000
BogoMIPS:                           5600.00
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          128 KiB
L1i cache:                          128 KiB
L2 cache:                           2 MiB
L3 cache:                           16 MiB
NUMA node0 CPU(s):                  0-7
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:             Mitigation; untrained return thunk; SMT enabled with STIBP protection
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; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected
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 tsc_known_freq pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr rdpru wbnoinvd arat npt nrip_save rdpid

Versions of relevant libraries:
[pip3] numpy==2.0.2
[pip3] torch==2.6.0.dev20241015+cpu
[pip3] torchvision==0.20.0.dev20241015+cpu
[conda] numpy                     2.0.2                    pypi_0    pypi
[conda] torch                     2.6.0.dev20241015+cpu          pypi_0    pypi
[conda] torchvision               0.20.0.dev20241015+cpu          pypi_0    pypi
```

cc @ezyang @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4