pytorch - 💡(How to fix) Fix SIGSEGV when assigning to nn.Parameter.data inside a torch.func transform

🐛 Describe the bug

When Parameter.data = <tensor> is executed inside a torch.func transform (e.g. torch.func.grad), subsequent use of that Parameter under the same transform crashes the process with SIGSEGV instead of running successfully or raising a catchable Python exception.

Reproduction:

import faulthandler, torch, torch.func as TF
faulthandler.enable()

class LinearFn(torch.autograd.Function):
    generate_vmap_rule = True
    @staticmethod
    def forward(input, w, b):
        return torch.addmm(b, input, w.t())
    @staticmethod
    def setup_context(ctx, inputs, output):
        ctx.save_for_backward(*inputs)
    @staticmethod
    def backward(ctx, gout):
        input, w, b = ctx.saved_tensors
        return gout.matmul(w), gout.t().matmul(input), gout.sum(0)

weight = torch.nn.Parameter(torch.randn(16, 16, device="cuda"))
bias = torch.nn.Parameter(torch.randn(16, device="cuda"))

def f(x):
    gathered_w = x.new_zeros(16, 16) + weight.detach()
    gathered_b = x.new_zeros(16) + bias.detach()
    weight.data = gathered_w
    bias.data = gathered_b
    return LinearFn.apply(x, weight, bias).sum()

x = torch.randn(4, 16, device="cuda")
TF.grad(f)(x)

Observed (exit code 139):

Fatal Python error: Segmentation fault

Current thread 0x00007ee8176dd140 (most recent call first):
  File ".../torch/utils/_pytree.py", line 1650 in wrapped
  File ".../torch/utils/_pytree.py", line 1280 in unflatten
  File ".../torch/utils/_pytree.py", line 1539 in tree_map
  File ".../torch/utils/_pytree.py", line 1715 in tree_map_only
  File ".../torch/_functorch/pyfunctorch.py", line 335 in dispatch_functorch
  File ".../torch/_ops.py", line 383 in dispatch
  File ".../torch/_ops.py", line 533 in __call__
  File ".../torch/_functorch/autograd_function.py", line 66 in __call__
  File ".../torch/autograd/function.py", line 606 in apply
  File "repro.py", line 25 in f
  File ".../torch/_functorch/eager_transforms.py", line 1457 in grad_and_value_impl
  File ".../torch/_functorch/vmap.py", line 56 in fn
  File ".../torch/_functorch/eager_transforms.py", line 1506 in grad_impl
  File ".../torch/_functorch/apis.py", line 433 in wrapper
  File "repro.py", line 29 in <module>
Segmentation fault (core dumped)

Expected: the program either runs to completion or raises a Python exception. SIGSEGV is never an acceptable response to user-level Python code.

Versions

PyTorch version: 2.11.0+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: Could not collect CMake version: Could not collect 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.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39 Is CUDA available: True CUDA runtime version: 13.0.48 CUDA_MODULE_LOADING set to: GPU models and configuration: GPU 0: NVIDIA GeForce RTX 5060 Ti GPU 1: NVIDIA GeForce RTX 5060 Ti

Nvidia driver version: 595.79 cuDNN version: Could not collect Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True Caching allocator config: N/A

CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 48 bits physical, 48 bits virtual Byte Order: Little Endian CPU(s): 16 Vendor ID: AuthenticAMD Model name: AMD Ryzen 9 9950X3D 16-Core Processor

Versions of relevant libraries: [pip3] numpy==2.4.3 [pip3] nvidia-cublas==13.1.0.3 [pip3] nvidia-cuda-cupti==13.0.85 [pip3] nvidia-cuda-nvrtc==13.0.88 [pip3] nvidia-cuda-runtime==13.0.96 [pip3] nvidia-cudnn-cu13==9.19.0.56 [pip3] nvidia-cufft==12.0.0.61 [pip3] nvidia-curand==10.4.0.35 [pip3] nvidia-cusolver==12.0.4.66 [pip3] nvidia-cusparse==12.6.3.3 [pip3] nvidia-cusparselt-cu13==0.8.0 [pip3] nvidia-nccl-cu13==2.28.9 [pip3] nvidia-nvjitlink==13.0.88 [pip3] nvidia-nvtx==13.0.85 [pip3] torch==2.11.0+cu130 [pip3] torchao==0.18.0+git357883dc [pip3] torchaudio==2.11.0+cu130 [pip3] torchvision==0.26.0+cu130 [pip3] triton==3.6.0 [conda] Could not collect

cc @chauhang @penguinwu @Chillee @samdow @kshitij12345