🐛[BUG]: SHTCUDA and iSHTCUDAisn't capturable with cuda graphs #11

import torch
from cuhpx import SHTCUDA, iSHTCUDA

nside = 32
lmax = 3*64 - 1
mmax = 3*64 - 1

sht = SHTCUDA(nside=nside, lmax=lmax, mmax=mmax, quad_weights='ring')
sht = sht.to(device="cuda")

x = torch.zeros((1,12,32,32)).to(device="cuda")

train_graph = torch.cuda.CUDAGraph()

capture_stream = torch.cuda.Stream()
capture_stream.wait_stream(torch.cuda.current_stream())

with torch.cuda.stream(capture_stream):
    for _ in range(5):
        out = sht(x)

    with torch.cuda.graph(train_graph):
        out = sht(x)

Traceback (most recent call last):
  File "/root/NV-dlesm-crps-loss/utils/./cuHPX_tester.py", line 23, in <module>
    out = sht(x)
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1510, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1519, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.12/dist-packages/cuhpx/hpx_sht.py", line 652, in forward
    return SHTFunction.apply(x, self.weights, self.pct, self.W, self.mmax, self.lmax, self.nside)
  File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 551, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/usr/local/lib/python3.12/dist-packages/cuhpx/hpx_sht.py", line 491, in forward
    x = cuhpx_fft.healpix_rfft_batch(x, mmax, nside)
RuntimeError: CUDA error: operation not permitted when stream is capturing
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.