added relevant headers for ROIAlign backwards

Author: varunagrawal

torchvision/csrc/ROIAlign.h

@@ -7,12 +7,13 @@
 #endif
 
 // Interface for Python
-at::Tensor ROIAlign_forward(const at::Tensor& input,
-                            const at::Tensor& rois,
-                            const float spatial_scale,
-                            const int pooled_height,
-                            const int pooled_width,
-                            const int sampling_ratio) {
+at::Tensor ROIAlign_forward(const at::Tensor& input,    // Input feature map.
+                            const at::Tensor& rois,     // List of ROIs to pool over.
+                            const float spatial_scale,  // The scale of the image features. ROIs will be scaled to this.
+                            const int pooled_height,    // The height of the pooled feature map.
+                            const int pooled_width,     // The width of the pooled feature
+                            const int sampling_ratio)   // The number of points to sample in each bin along each axis.
+{                   
   if (input.type().is_cuda()) {
 #ifdef WITH_CUDA
     return ROIAlign_forward_cuda(input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio);
@@ -40,6 +41,6 @@ at::Tensor ROIAlign_backward(const at::Tensor& grad,
     AT_ERROR("Not compiled with GPU support");
 #endif
   }
-  AT_ERROR("Not implemented on the CPU");
+  return ROIAlign_backward_cpu(grad, rois, spatial_scale, pooled_height, pooled_width, batch_size, channels, height, width, sampling_ratio);
 }
 

torchvision/csrc/cpu/vision.h

@@ -25,6 +25,17 @@ at::Tensor ROIAlign_forward_cpu(const at::Tensor &input,
                                 const int pooled_width,
                                 const int sampling_ratio);
 
+at::Tensor ROIAlign_backward_cpu(const at::Tensor &grad,
+                                 const at::Tensor &rois,
+                                 const float spatial_scale,
+                                 const int pooled_height,
+                                 const int pooled_width,
+                                 const int batch_size,
+                                 const int channels,
+                                 const int height,
+                                 const int width,
+                                 const int sampling_ratio);
+
 at::Tensor nms_cpu(const at::Tensor &dets,
                    const at::Tensor &scores,
                    const float threshold);