Repository: pytorch
Issue: __torch_dispatch__ fails on functionalization when lazy device and `out=` arg is used
Labels: triaged,module: __torch_dispatch__,module: lazy,module: functionalization

### 🐛 Describe the bug

__torch_dispatch__ leads to:
```
RuntimeError: false INTERNAL ASSERT FAILED at "aten/src/ATen/RegisterFunctionalization_0.cpp":2151, please report a bug to PyTorch. mutating a non-functional tensor with a functional tensor is not allowed. Please ensure that all of your inputs are wrapped inside of a functionalize() call.
```
when lazy device and `out=` arg is used.

Example:
```py
import torch
import torch._lazy

import torch._lazy.ts_backend
torch._lazy.ts_backend.init()

from torch.utils._python_dispatch import TorchDispatchMode

class CustomDispatchMode(TorchDispatchMode):
    def __torch_dispatch__(self, func, types, args, kwargs):
        return func(*args, **kwargs)

a = torch.ones((16), dtype=torch.float, device='lazy')
b = torch.ones((16), dtype=torch.float, device='lazy')

with CustomDispatchMode():
    out = torch.empty(16, dtype=a.dtype).to(device="lazy")
    torch.add(a, b, out=out)
```

The same example works fine for the CPU device.

CC: @albanD @bdhirsh 

### Versions

```
Collecting environment information...
PyTorch version: 2.1.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 7.9.2009 (Core) (x86_64)
GCC version: Could not collect
Clang version: 3.4.2 (tags/RELEASE_34/dot2-final)
CMake version: version 3.27.1
Libc version: glibc-2.30

Python version: 3.8.16 (default, Nov  2 2023, 16:16:22)  [GCC 8.4.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.80.1.el7.x86_64-x86_64-with-glibc2.2.5
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
Byte Order:            Little Endian
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:                 1
Model name:            AMD EPYC 7571
Stepping:              2
CPU MHz:               2199.706
BogoMIPS:              4399.41
Hypervisor vendor:     KVM
Virtualization type:   full
L1d cache:             32K
L1i cache:             64K
L2 cache:              512K
L3 cache:              8192K
NUMA node0 CPU(s):     0-7
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 art rep_good nopl nonstop_tsc extd_apicid aperfmperf eagerfpu 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 perfctr_core rsb_ctxsw vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt sha_ni xsaveopt xsavec xgetbv1 retpoline_amd clzero xsaveerptr arat npt nrip_save

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] torch==2.1.1+cpu
[pip3] torchvision==0.16.1+cpu
[conda] Could not collect
```

cc @Chillee @ezyang @zou3519 @albanD @samdow @bdhirsh