pytorch - 💡(How to fix) Fix `torch.compile` + `torch.while_loop`: stride mismatch crash during backward pass

InductorError: UncapturedHigherOrderOpError: Expected carried_inputs and body_output
to have same metadata but found:
pair[1] differ in 'stride: (0, 0) vs (64, 1)', where lhs is
FakeTensor(..., device='cuda:0', size=(4, 64)) and rhs is
FakeTensor(..., device='cuda:0', size=(4, 64))

While executing %while_loop : [num_users=1] = call_function[target=torch.ops.higher_order.while_loop](...)

---

import torch

torch.manual_seed(42)
x = torch.randn(4, 64, device='cuda', requires_grad=True)

def cond_fn(i, x):
    return i < 5

def body_fn(i, x):
    return i + 1, x * 0.9 + 0.1

def fn(x):
    _, result = torch.while_loop(
        cond_fn, body_fn,
        (torch.tensor(0, device='cuda'), x)
    )
    return result.sum()

# Eager — OK
loss = fn(x)
loss.backward()
print(f"Eager OK: grad_max={x.grad.abs().max().item():.6f}")
x.grad = None

# Compiled — CRASH
torch._dynamo.reset()
fn_c = torch.compile(fn, fullgraph=True)
loss_c = fn_c(x)
loss_c.backward()  # <-- crashes here

---

PyTorch: 2.13.0.dev20260501+cu126
Python: 3.11
CUDA: 12.6
GPU: Tesla T4