How to use the torchkbnufft.functional.kbinterp.KbInterpFunction.apply function in torchkbnufft

Args:
            x (tensor): The DFT of the signal.
            om (tensor, optional): A new set of omega coordinates to
                interpolate to in radians/voxel.
            interp_mats (dict, default=None): A dictionary with keys
                'real_interp_mats' and 'imag_interp_mats', each key containing
                a list of interpolation matrices (see 
                mri.sparse_interp_mat.precomp_sparse_mats for construction).
                If None, then a standard interpolation is run.

        Returns:
            tensor: x computed at off-grid locations in om.
        """
        interpob = self._extract_interpob()

        y = KbInterpFunction.apply(x, om, interpob, interp_mats)

        return y

This function wraps scale_and_fft_on_image_volume and KbInterpFunction
        for PyTorch autograd.
        """
        x = x.clone()

        # extract interpolation params
        scaling_coef = interpob['scaling_coef']
        grid_size = interpob['grid_size']
        im_size = interpob['im_size']
        norm = interpob['norm']

        x = scale_and_fft_on_image_volume(
            x, scaling_coef, grid_size, im_size, norm)

        y = KbInterpFunction.apply(x, om, interpob, interp_mats)

        ctx.save_for_backward(om)
        ctx.interpob = interpob
        ctx.interp_mats = interp_mats

        return y

"""
        x = x.clone()

        om, = ctx.saved_tensors
        interpob = ctx.interpob
        interp_mats = ctx.interp_mats

        scaling_coef = interpob['scaling_coef']
        grid_size = interpob['grid_size']
        im_size = interpob['im_size']
        norm = interpob['norm']

        x = scale_and_fft_on_image_volume(
            x, scaling_coef, grid_size, im_size, norm)

        y = KbInterpFunction.apply(x, om, interpob, interp_mats)

        return y, None, None, None