Minor cleanup of roi_align_forward_kernel_impl (#3619)

* minor clean up

* do same for ps_roialign

Co-authored-by: Francisco Massa <fvsmassa@gmail.com>

Author: NicolasHug

torchvision/csrc/ops/cpu/ps_roi_align_kernel.cpp

@@ -62,7 +62,7 @@ T bilinear_interpolate(
 
 template <typename T>
 void ps_roi_align_forward_kernel_impl(
-    int nthreads,
+    int num_rois,
     const T* input,
     const T spatial_scale,
     int channels,
@@ -75,7 +75,6 @@ void ps_roi_align_forward_kernel_impl(
     int channels_out,
     T* output,
     int* channel_mapping) {
-  int num_rois = nthreads / channels_out / pooled_width / pooled_height;
   for (int n = 0; n < num_rois; n++) {
     // [start, end) interval for spatial sampling
     const T* offset_rois = rois + n * 5;
@@ -335,16 +334,15 @@ std::tuple<at::Tensor, at::Tensor> ps_roi_align_forward_kernel(
   auto channel_mapping =
       at::zeros(output.sizes(), input.options().dtype(at::kInt));
 
-  auto output_size = output.numel();
-  if (output_size == 0) {
+  if (output.numel() == 0) {
     return std::make_tuple(output, channel_mapping);
   }
 
   auto input_ = input.contiguous(), rois_ = rois.contiguous();
   AT_DISPATCH_FLOATING_TYPES_AND_HALF(
       input.scalar_type(), "ps_roi_align_forward_kernel", [&] {
         ps_roi_align_forward_kernel_impl<scalar_t>(
-            output_size,
+            num_rois,
             input_.data_ptr<scalar_t>(),
             spatial_scale,
             channels,

torchvision/csrc/ops/cpu/roi_align_kernel.cpp

@@ -117,7 +117,7 @@ void pre_calc_for_bilinear_interpolate(
 
 template <typename T>
 void roi_align_forward_kernel_impl(
-    int nthreads,
+    int n_rois,
     const T* input,
     const T& spatial_scale,
     int channels,
@@ -129,7 +129,6 @@ void roi_align_forward_kernel_impl(
     bool aligned,
     const T* rois,
     T* output) {
-  int n_rois = nthreads / channels / pooled_width / pooled_height;
   // (n, c, ph, pw) is an element in the pooled output
   // can be parallelized using omp
   // #pragma omp parallel for num_threads(32)
@@ -414,16 +413,14 @@ at::Tensor roi_align_forward_kernel(
   at::Tensor output = at::zeros(
       {num_rois, channels, pooled_height, pooled_width}, input.options());
 
-  auto output_size = num_rois * pooled_height * pooled_width * channels;
-
   if (output.numel() == 0)
     return output;
 
   auto input_ = input.contiguous(), rois_ = rois.contiguous();
   AT_DISPATCH_FLOATING_TYPES_AND_HALF(
       input.scalar_type(), "roi_align_forward_kernel", [&] {
         roi_align_forward_kernel_impl<scalar_t>(
-            output_size,
+            num_rois,
             input_.data_ptr<scalar_t>(),
             spatial_scale,
             channels,