content: import torch
  from torch import nn
  from torch.autograd import Function
  from torch.autograd.function import once_differentiable
  from torch.nn.modules.utils import _pair
  
  # import alphaction._custom_cuda_ext as _C
  from torchvision.ops import roi_align
  
  
  # class _ROIAlign3d(Function):
  #     @staticmethod
  #     def forward(ctx, input, roi, output_size, spatial_scale, sampling_ratio):
  #         ctx.save_for_backward(roi)
  #         ctx.output_size = _pair(output_size)
  #         ctx.spatial_scale = spatial_scale
  #         ctx.sampling_ratio = sampling_ratio
  #         ctx.input_shape = input.size()
  #         output = _C.roi_align_3d_forward(
  #             input, roi, spatial_scale, output_size[0], output_size[1], sampling_ratio
  #         )
  #         return output
  #
  #     @staticmethod
  #     @once_differentiable
  #     def backward(ctx, grad_output):
  #         rois, = ctx.saved_tensors
  #         output_size = ctx.output_size
  #         spatial_scale = ctx.spatial_scale
  #         sampling_ratio = ctx.sampling_ratio
  #         bs, ch, l, h, w = ctx.input_shape
  #         grad_input = _C.roi_align_3d_backward(
  #             grad_output,
  #             rois,
  #             spatial_scale,
  #             output_size[0],
  #             output_size[1],
  #             bs,
  #             ch,
  #             l,
  #             h,
  #             w,
  #             sampling_ratio,
  #         )
  #         return grad_input, None, None, None, None
  #
  #
  # roi_align_3d = _ROIAlign3d.apply
  
  
  class ROIAlign3d(nn.Module):
      def __init__(self, output_size, spatial_scale, sampling_ratio):
          super(ROIAlign3d, self).__init__()
          self.output_size = output_size
          self.spatial_scale = spatial_scale
          self.sampling_ratio = sampling_ratio
  
      def forward(self, input, rois):
          # return roi_align_3d(
          #     input, rois, self.output_size, self.spatial_scale, self.sampling_ratio
          # )
          # aia提供的algin不支持半精度，先使用torchvision
          return roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio)
  
      def __repr__(self):
          tmpstr = self.__class__.__name__ + "("
          tmpstr += "output_size=" + str(self.output_size)
          tmpstr += ", spatial_scale=" + str(self.spatial_scale)
          tmpstr += ", sampling_ratio=" + str(self.sampling_ratio)
          tmpstr += ")"
          return tmpstr
  