pytorch - 💡(How to fix) Fix "ZeroTensors are immutable" error while optimizing partial derivatives [1 comments, 2 participants]

🐛 Describe the bug

Trying to compute a backward pass on the partial derivatives (computed via torch.func.jvp) of a compiled network fails with a "ZeroTensors are immutable" error. If I disable torch.compile, skip partial differentiation, or simplify the network, the error goes away. Forward computation of the same function seems to run fine. Some specific sizes of the network cause an CUDA illegal memory access error instead.

This example reproduces the issue:

import torch

def jvp(f, x):
    return torch.func.jvp(f, (x.clone(),), (torch.ones_like(x),))[1] # replace [1] with [0] to use values instead of derivatives => runs (but doesn't compute what we want)

@torch.compile() # remove this line to disable compilation => runs (but slow)
def compute(f, x):
    first, rest = x[..., :1], x[..., 1:]
    return jvp(lambda X: f(torch.cat([X, rest])), first) # return partial derivative w.r.t. first dimension

def Simple(in_features, hidden_features, out_features):
    return torch.nn.Sequential(
            torch.nn.Linear(in_features, hidden_features),
            torch.nn.Linear(hidden_features, hidden_features), # remove this line to simplify the network => runs (but doesn't match the model that we'd like to use)
            torch.nn.Linear(hidden_features, out_features))

in_features = 4 # set this to 1 => CUDA illegal memory access
net = Simple(in_features, 32, 8).cuda()
opt = torch.optim.SGD(params = net.parameters())

x = torch.rand((in_features,)).cuda()
L = compute(net, x).sum()

print(L.item()) # this prints so forward pass should be OK
L.backward() # errors trigger on backward

Error logs

RuntimeError: ZeroTensors are immutable. Please use the materialized zero tensor obtained using .clone() if you want a mutable tensor.

Versions

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

OS: Ubuntu 22.04.2 LTS (x86_64) GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04.3) 11.4.0 Clang version: Could not collect CMake version: Could not collect Libc version: glibc-2.35

Python version: 3.11.14 (main, Feb 12 2026, 00:42:50) [Clang 21.1.4 ] (64-bit runtime) Python platform: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.35 Is CUDA available: True CUDA runtime version: Could not collect CUDA_MODULE_LOADING set to: GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090 Nvidia driver version: 595.71 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): 24 On-line CPU(s) list: 0-23 Vendor ID: AuthenticAMD Model name: AMD Ryzen 9 5900X 12-Core Processor CPU family: 25 Model: 33 Thread(s) per core: 2 Core(s) per socket: 12 Socket(s): 1 Stepping: 0 BogoMIPS: 7399.98 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 rep_good nopl cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves clzero xsaveerptr arat npt nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload umip vaes vpclmulqdq rdpid fsrm Virtualization: AMD-V Hypervisor vendor: Microsoft Virtualization type: full L1d cache: 384 KiB (12 instances) L1i cache: 384 KiB (12 instances) L2 cache: 6 MiB (12 instances) L3 cache: 32 MiB (1 instance) NUMA node(s): 1 NUMA node0 CPU(s): 0-23 Vulnerability Gather data sampling: Not affected Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Not affected Vulnerability Mds: Not affected Vulnerability Meltdown: Not affected Vulnerability Mmio stale data: Not affected Vulnerability Reg file data sampling: Not affected Vulnerability Retbleed: Not affected Vulnerability Spec rstack overflow: Vulnerable: Safe RET, no microcode Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization Vulnerability Spectre v2: Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP conditional; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected Vulnerability Srbds: Not affected Vulnerability Tsx async abort: Not affected

Versions of relevant libraries: [pip3] Could not collect [conda] Could not collect

cc @ezyang @gchanan @kadeng @msaroufim @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @muchulee8 @amjames @aakhundov @coconutruben @jataylo @albanD @gqchen @nikitaved @soulitzer @Varal7 @bobrenjc93