[BUG] Possible memory leak in backward pass of custom torch operator when used with torch DDP.

Bug Description

I have a custom pytorch operator integrated with a warp kernel that I coded up following the documentation here, https://nvidia.github.io/warp/modules/interoperability.html#example-optimization-using-pytorch-custom-operators-pytorch-2-4-0.

The operator is designed to have a handful of trainable parameters associated with it. When I run a training loop with this operator with optimization of these parameters turned off GPU memory usage stays low and constant.

Running the same training loop with optimization of the parameters turned on using a single GPU and no DDP gives a cycle of memory increases and apparent garbage collection but still the total usage appears bounded.

Running with multiple GPUs (8 in this case) gives the same GC cycle behavior but it looks like there are errors and not all the memory is freed. Eventually these GC cycles stop happening and the process eventually goes out of memory. Also the logs for these runs throw a bunch of Cuda errors that appear to be coincident with an attempt to perform GC.

`
import torch
from models.WarpSetConvRBF import WarpKernels as WPRBF
import warp as wp

@torch.library.custom_op("wp::warp_2dRBKF", mutates_args=())
def warp_2dRBKF(sigma:torch.Tensor,
phi: torch.Tensor,
x: torch.Tensor,
y: torch.Tensor,
x_hat: torch.Tensor,
y_hat: torch.Tensor,
deltax: float,
deltay: float) -> torch.Tensor:
"""Wrap WarpKernels.sum_rbkf_3 as a torch operator

Args:
    phi (torch.Tensor): Array of lifted observed values [degree+1, nobs].
    x (torch.Tensor): x-coordinates of grid [N_x]
    deltax (float):  Distance between x-coordinates of grid x[1] - x[0]
    y (torch.Tensor): y-coordinates of grid [N_y]
    deltay (float): Distance between y-coordinates of grid y[1] - y[0]
    x_hat (torch.Tensor): x-coordinates of observations [nobs].
    y_hat (torch.Tensor): y-coordinates of observations [nobs].
    sigma (torch.Tensor): Width of kernels sigma = [sigma_x, sigma_y]

Returns:
    torch.Tensor: 
        $$
        E(X,Y,T | \hat{X}, \hat{Y}) = \sum_{i=1}^NObs phi_i exp(-(x - \hat{x}_i)^2/\sigma_x^2)exp(-(y - \hat{y}_i)^2/\sigma_y^2).
        $$
"""
wp_phi = wp.from_torch(phi,dtype=wp.float32,return_ctype=True)
[degp1,nobs] = phi.shape
wp_x = wp.from_torch(x, dtype=wp.float32,return_ctype=True)
wp_y = wp.from_torch(y, dtype=wp.float32,return_ctype=True)
wp_x_hat = wp.from_torch(x_hat, dtype=wp.float32,return_ctype=True)
wp_y_hat = wp.from_torch(y_hat, dtype=wp.float32,return_ctype=True)
wp_sigma = wp.from_torch(sigma, dtype=wp.float32,return_ctype=True)

wp_ret = wp.zeros([degp1,len(y),len(x)], dtype=wp.float32)

wp.launch(kernel=WPRBF.sum_rbkf_2, dim=[degp1,nobs],
          inputs=[
              wp_sigma,
              wp_phi,
              wp_x,
              wp_y,
              wp_x_hat,
              wp_y_hat,
              deltax,
              deltay
          ],
          outputs=[wp_ret])

return wp.to_torch(wp_ret)

@warp_2dRBKF.register_fake
def _(sigma: torch.Tensor,
phi: torch.Tensor,
x: torch.Tensor,
y: torch.Tensor,
x_hat: torch.Tensor,
y_hat: torch.Tensor,
deltax: float,
deltay: float):
[degp1, _] = phi.shape
return torch.empty([degp1,len(y),len(x)], dtype=torch.float32)

@torch.library.custom_op("wp::warp_2dRBKF_backward", mutates_args=())
def warp_2dRBKF_backward(
sigma: torch.Tensor,
phi: torch.Tensor,
x: torch.Tensor,
y: torch.Tensor,
x_hat: torch.Tensor,
y_hat: torch.Tensor,
deltax: float,
deltay: float,
f: torch.Tensor,
adj_f: torch.Tensor
) -> torch.Tensor:
"""Wrap adjoint of WarpKernels.sum_rbkf_2 as a torch operator

Args:
    phi (torch.Tensor): Array of lifted observed values [degree+1, nobs].
    x (torch.Tensor): x-coordinates of grid [N_x]
    deltax (float):  Distance between x-coordinates of grid x[1] - x[0]
    y (torch.Tensor): y-coordinates of grid [N_y]
    deltay (float): Distance between y-coordinates of grid y[1] - y[0]
    x_hat (torch.Tensor): x-coordinates of observations [nobs].
    y_hat (torch.Tensor): y-coordinates of observations [nobs].
    sigma (torch.Tensor): Width of kernels sigma = [sigma_x, sigma_y]
    f (torch.Tensor): Forward pass output
    adj_f (torch.Tensor): Adjoint output

Returns:
    torch.Tensor: gradient with respect to sigma
"""
wp_phi = wp.from_torch(phi,dtype=wp.float32,return_ctype=True)
[degp1,nobs] = phi.shape
wp_x = wp.from_torch(x, dtype=wp.float32,return_ctype=True)
wp_y = wp.from_torch(y, dtype=wp.float32,return_ctype=True)
wp_x_hat = wp.from_torch(x_hat, dtype=wp.float32,return_ctype=True)
wp_y_hat = wp.from_torch(y_hat, dtype=wp.float32,return_ctype=True)
wp_sigma = wp.from_torch(sigma, dtype=wp.float32)

wp_f = wp.from_torch(f, requires_grad=True,return_ctype=True)
wp_adj_f = wp.from_torch(adj_f, requires_grad=False,return_ctype=True)

wp.launch(
    kernel=WPRBF.sum_rbkf_2,
    dim=[degp1,nobs],
    inputs=[
              wp_sigma,
              wp_phi,
              wp_x,
              wp_y,
              wp_x_hat,
              wp_y_hat,
              deltax,
              deltay
      ],
    outputs=[wp_f],
    adj_inputs=[wp_sigma.grad,None,None,None,None,None,None,None],
    adj_outputs=[wp_adj_f],
    adjoint=True
)

return wp.to_torch(wp_sigma.grad)

@warp_2dRBKF_backward.register_fake
def _( sigma: torch.Tensor,
phi: torch.Tensor,
x: torch.Tensor,
y: torch.Tensor,
x_hat: torch.Tensor,
y_hat: torch.Tensor,
deltax: float,
deltay: float,
f: torch.Tensor,
adj_f: torch.Tensor):
return torch.empty_like(sigma)

def backward(ctx, adj_f):
"""Implements autograd for RBKF

Args:
    ctx : I don't really understand this code. It's more or less copied from here
       https://nvidia.github.io/warp/modules/interoperability.html#example-optimization-using-pytorch-custom-operators-pytorch-2-4-0
    adj_f : _description_

Returns:
    _type_: _description_
"""
ctx.sigma.grad = warp_2dRBKF_backward(ctx.sigma,ctx.phi,
              ctx.x,
              ctx.y,
              ctx.x_hat,
              ctx.y_hat,
              ctx.deltax,
              ctx.deltay,ctx.f, adj_f)
return ctx.sigma.grad, None, None, None, None, None, None, None

def setup_context(ctx, inputs, output):
ctx.sigma,ctx.phi,ctx.x, ctx.y, ctx.x_hat, ctx.y_hat,ctx.deltax,ctx.deltay = inputs
ctx.f = output

warp_2dRBKF.register_autograd(backward, setup_context=setup_context)
`

System Information

OS: Ubuntu 22.04.5
Python 3.11.10
pyTorch: '2.7.0+cu126'
warp: 1.7.1