import torch

class TwoLayerMLP(torch.nn.Module):
    def __init__(self, d):
        super().__init__()
        self.fc1 = torch.nn.Linear(d, d)
        self.fc2 = torch.nn.Linear(d, d)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

d = 8
module = TwoLayerMLP(d)                           # parameters on CPU (default)
example_input = torch.randn(1, d, device="cuda")  # input on CUDA

print(f"module params device: {next(module.parameters()).device}")
print(f"example input device: {example_input.device}")

try:
    ep = torch.export.export(module, (example_input,))
    print(f"exported: {type(ep).__name__}")
except Exception as e:
    print(f"{type(e).__name__}: {e}")

# Reference: CPU export works fine
ep_cpu = torch.export.export(TwoLayerMLP(d), (torch.randn(1, d),))
print(f"cpu export: ok ({type(ep_cpu).__name__})")

---

module params device: cpu
example input device: cuda:0
RuntimeError: Unhandled FakeTensor Device Propagation for aten.mm.default,
  found two different devices cpu, cuda:0
cpu export: ok (ExportedProgram)

---

module params device: cpu
example input device: cuda:0
ValueError: Module parameters are on cpu but example inputs are on cuda:0.
  Move the module to the target device before exporting:
  module = module.to('cuda') and re-run torch.export.export.
cpu export: ok (ExportedProgram)

---

# MPS repro
module = torch.nn.Linear(4, 4)                    # params on CPU
example_input = torch.randn(4, 4, device="mps")
torch.export.export(module, (example_input,))
# RuntimeError: Unhandled FakeTensor Device Propagation for aten.mm.default,
#   found two different devices cpu, mps:0

---

PyTorch version: 2.13.0.dev20260512+cu130
Is debug build: False
CUDA used to build PyTorch: 13.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0
Clang version: 18.1.3 (1ubuntu1)
CMake version: version 3.28.3
Libc version: glibc-2.39

Python version: 3.12.3 (main, Mar 23 2026, 19:04:32) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.14.0-37-generic-x86_64-with-glibc2.39
Is CUDA available: True
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 5090
Nvidia driver version: 590.48.01

[pip3] numpy==2.4.4
[pip3] torch==2.13.0.dev20260512+cu130
[pip3] triton==3.7.0+git88b227e2