Repository: pytorch
Issue: torch.export(strict=False) fails with tensor constants
Labels: triaged,oncall: export

### 🐛 Describe the bug

Repro:

```
class Model(torch.nn.Module):
    def __init__(self, device):
        super().__init__()
        self.b = torch.randn(4, device=device)

    def forward(self, x):
        return x + self.b

device = torch.device("cpu")
sample_inputs = (torch.randn(4, 4, device=device),)

exported_program = torch.export.export(
    Model(device),
    sample_inputs,
    strict=False,
)
``` 

```

  File "/usr/local/fbcode/platform010/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/scripts/bahuang/export_torchbind.py", line 28, in <module>
    Model(device),
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/export/__init__.py", line 186, in export
    return _export(
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/export/exported_program.py", line 81, in wrapper
    return fn(*args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/export/_trace.py", line 700, in _export
    ep_non_strict = _export_non_strict(
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/export/_trace.py", line 421, in _export_non_strict
    gm, graph_signature = transform(aot_export_module)(
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/export/_trace.py", line 670, in _aot_export_non_strict
    gm, sig = aot_export(wrapped_mod, args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_functorch/aot_autograd.py", line 1032, in aot_export_module
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_functorch/aot_autograd.py", line 1217, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_dynamo/utils.py", line 248, in time_wrapper
    r = func(*args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_functorch/aot_autograd.py", line 509, in create_aot_dispatcher_function
    fw_metadata = run_functionalized_fw_and_collect_metadata(
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_functorch/_aot_autograd/collect_metadata_analysis.py", line 113, in inner
    flat_f_outs = f(*flat_f_args)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_functorch/_aot_autograd/utils.py", line 161, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 622, in functional_call
    out = mod(*args[params_len:], **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/export/_trace.py", line 657, in forward
    tree_out = self._export_root(*args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/scripts/bahuang/export_torchbind.py", line 22, in forward
    return x + self.b
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/export/_safeguard.py", line 42, in __torch_function__
    return func(*args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_subclasses/functional_tensor.py", line 380, in __torch_dispatch__
    outs_unwrapped = func._op_dk(
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_subclasses/fake_tensor.py", line 1637, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_subclasses/fake_tensor.py", line 1975, in dispatch
    return self._dispatch_impl(func, types, args, kwargs)
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_subclasses/fake_tensor.py", line 2050, in _dispatch_impl
    (flat_args, flat_arg_fake_tensors) = self.validate_and_convert_non_fake_tensors(
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_subclasses/fake_tensor.py", line 2281, in validate_and_convert_non_fake_tensors
    validated_args = [validate(a) for a in flat_args]
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_subclasses/fake_tensor.py", line 2281, in <listcomp>
    validated_args = [validate(a) for a in flat_args]
  File "/data/users/bahuang/fbsource/buck-out/v2/gen/fbcode/6f27a84d3075b0d5/scripts/bahuang/__export_torchbind__/export_torchbind#link-tree/torch/_subclasses/fake_tensor.py", line 2271, in validate
    raise Exception(
Exception: Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in aten.add.Tensor(FakeTensor(..., size=(4, 4)), tensor([...], size=(4,)))
``` 




### Versions

master

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