[Bug] @torch.jit.script + CUDA graph causes segfault on A30 (CUDA 12.8)

[alexhegit]

Bug Description

Using torch.jit.script decorated functions alongside CUDA graph replay causes a segmentation fault on NVIDIA A30 GPUs with CUDA 12.8. This appears to be an interaction issue between PyTorch's JIT scripting and CUDA graph capture.

Steps to Reproduce

1. Set up environment: A30 GPU + CUDA 12.8 + PyTorch with CUDA support
2. Run any training script that uses JIT-scripted functions with CUDA graphs enabled (default behavior)
3. Training crashes with segfault during CUDA graph capture/replay

Affected files:

• protomotions/kernels/rotations.py — 4 occurrences of @torch.jit.script
• protomotions/utils/task.py — 3 occurrences of @torch.jit.script

Expected Behavior

Training should run without segfault on all supported GPU configurations.

Actual Behavior

Process terminates with SIGSEGV during CUDA graph operations. The crash is deterministic when the JIT-compiled functions are invoked during graph capture.

Mitigation

Temporarily removing @torch.jit.script decorators from the affected functions allows training to complete on A30 + CUDA 12.8 without crashing.

Environment

• GPU: NVIDIA A30 (Compute Capability 8.0)
• CUDA: 12.8
• PyTorch: latest stable with CUDA 12.8 support
• OS: Linux

Questions for Maintainers

1. Is @torch.jit.script required for correct functionality, or is it primarily a performance optimization?
2. Would a version guard (e.g., checking torch.version.cuda and GPU architecture) be appropriate to conditionally disable JIT scripting on affected configurations?
3. Are there known incompatibilities between torch.jit.script and CUDA graph replay in certain PyTorch/CUDA versions combinations?

Additional Notes

• This issue does not occur on all GPU/CUDA combinations — it is specific to the A30 + CUDA 12.8 configuration tested.
• The issue appears to be an environmental incompatibility rather than a logical bug in the code itself.
• A fix is not urgently needed for all users, but maintainers may want to be aware of this interaction for future PyTorch/CUDA compatibility.

[tesslerc]

Thank you for the report.
Have you been able to resolve it? For example, using some other combination of packages?