Repository: pytorch
Issue: `grad_fn` is not defined in `__torch_dispatch__` arguments when lazy device is used
Labels: triaged,lazy

### 🐛 Describe the bug

I’m trying to use TorchDispatchMode to hook grad_fn for every tensor created within a dispatch mode context. I can successfully access grad_fn when running on a CPU device. However, when I’m switching to a lazy device, grad_fn is not defined for all tensors inside TestDispatchMode context.

This is the example I’m using to debug __torch_dispatch__ behaviour:
```py
# This is a simple dispatch mode that prints grad_fn 
# for every input tensor passed to __torch_dispatch__.
class TestDispatchMode(TorchDispatchMode):
    def __torch_dispatch__(self, func, types, args, kwargs):
        def check_grad_fn(tensor):
            if isinstance(tensor, torch.Tensor):
                print(f"    {tensor.grad_fn=}")
        
        print(f"{func=}")

        tree_map(check_grad_fn, args)
        tree_map(check_grad_fn, kwargs)

        return func(*args, **kwargs)

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(5, 5, bias=False)

    def forward(self, x):
        out = self.fc1(x)
        assert out.grad_fn # <--- grad_fn is defined here for both CPU and lazy devices

        with TestDispatchMode():
            # grad_fn is defined for `out` tensor in torch_dispatch for CPU device,
            # however it is none for the lazy device.
            out = F.relu(out)

        return out
```
So, if I run this model using a CPU device
```py
model = Model()
model.train()
out = model(torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.float32))
```
I get the expected output where each input tensor has grad_fn
```
    """
    output:
    func=<OpOverload(op='aten.relu', overload='default')>
        tensor.grad_fn=<MmBackward0 object at 0x7f7a309aeaf0>
    func=<OpOverload(op='aten.detach', overload='default')>
        tensor.grad_fn=<ReluBackward0 object at 0x7f7a309aeaf0>
    """
```
however,  if I switch to the lazy device
```py
import torch._lazy.ts_backend
torch._lazy.ts_backend.init()

model = Model().to('lazy')
model.train()
out = model(torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.float32, device='lazy'))
```
and run the same thing using the TS backend I’ll get nones:
```
    """
    output:
    func=<OpOverload(op='aten.relu', overload='default')>
        tensor.grad_fn=None
    """
```
CC: @JackCaoG @bdhirsh 

### Versions

```
PyTorch version: 2.1.0+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, Oct 10 2023, 22:55:17)  [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):                2
On-line CPU(s) list:   0,1
Thread(s) per core:    2
Core(s) per socket:    1
Socket(s):             1
NUMA node(s):          1
Vendor ID:             GenuineIntel
CPU family:            6
Model:                 85
Model name:            Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
Stepping:              7
CPU MHz:               2499.998
BogoMIPS:              4999.99
Hypervisor vendor:     KVM
Virtualization type:   full
L1d cache:             32K
L1i cache:             32K
L2 cache:              1024K
L3 cache:              36608K
NUMA node0 CPU(s):     0,1
Flags:                 fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc aperfmperf eagerfpu pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single rsb_ctxsw fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 ida arat pku ospke

Versions of relevant libraries:
[pip3] numpy==1.18.2
[pip3] torch==2.1.0+cpu
[pip3] torchdata==0.3.0
[pip3] torchtext==0.12.0
[pip3] torchvision==0.16.0+cpu
```