support for Half types in ROIAlign

Author: varunagrawal

torchvision/csrc/cpu/ROIAlign_cpu.cpp

@@ -369,7 +369,7 @@ at::Tensor ROIAlign_forward_cpu(const at::Tensor& input,
   if (output.numel() == 0)
     return output;
 
-  AT_DISPATCH_FLOATING_TYPES(input.type(), "ROIAlign_forward", [&] {
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "ROIAlign_forward", [&] {
     ROIAlignForward<scalar_t>(
          output_size,
          input.data<scalar_t>(),
@@ -414,7 +414,7 @@ at::Tensor ROIAlign_backward_cpu(const at::Tensor& grad,
   int h_stride = grad.stride(2);
   int w_stride = grad.stride(3);
 
-  AT_DISPATCH_FLOATING_TYPES(grad.type(), "ROIAlign_forward", [&] {
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad.type(), "ROIAlign_forward", [&] {
     ROIAlignBackward<scalar_t>(
         grad.numel(),
         grad.data<scalar_t>(),

torchvision/csrc/cuda/ROIAlign_cuda.cu

@@ -273,7 +273,7 @@ at::Tensor ROIAlign_forward_cuda(const at::Tensor& input,
     return output;
   }
 
-  AT_DISPATCH_FLOATING_TYPES(input.type(), "ROIAlign_forward", [&] {
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "ROIAlign_forward", [&] {
     RoIAlignForward<scalar_t><<<grid, block, 0, stream>>>(
          output_size,
          input.data<scalar_t>(),
@@ -323,7 +323,7 @@ at::Tensor ROIAlign_backward_cuda(const at::Tensor& grad,
   int h_stride = grad.stride(2);
   int w_stride = grad.stride(3);
 
-  AT_DISPATCH_FLOATING_TYPES(grad.type(), "ROIAlign_backward", [&] {
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad.type(), "ROIAlign_backward", [&] {
     RoIAlignBackward<scalar_t><<<grid, block, 0, stream>>>(
          grad.numel(),
          grad.data<scalar_t>(),