commit 24601daa1203a9ad1232e1d18a07ff4842d53d27
author Zafar Takhirov <cc.rafaz@zafar.cc> Tue Jul 23 20:56:35 2019 -0700
committer Facebook Github Bot <facebook-github-bot@users.noreply.github.com> Tue Jul 23 21:11:06 2019 -0700
tree 09927bf8d105e5b137278ea61659870d986a3cb2
parent e7a9b0d62f6ca6f3e9f350e827ff7b58a2cb7bdb

Adding check for a single batch in adaptive_avg_pool

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/23137

Test Plan: Imported from OSS

Differential Revision: D16403804

Pulled By: zafartahirov

fbshipit-source-id: df79a8c768ffabeceb4c0044c967a623c5885484

aten/src/ATen/native/AdaptiveAveragePooling.cpp

diff --git a/aten/src/ATen/native/AdaptiveAveragePooling.cpp b/aten/src/ATen/native/AdaptiveAveragePooling.cpp
index eec9c75..616cda3 100644
--- a/aten/src/ATen/native/AdaptiveAveragePooling.cpp
+++ b/aten/src/ATen/native/AdaptiveAveragePooling.cpp
@@ -129,10 +129,13 @@
     auto osizeW = output_size[1];
 
     /* resize output */
-    if (input.ndimension() == 3)
+    if (input.ndimension() == 3 || input.size(-4) == 1)
     {
-      output.resize_({sizeD, osizeH, osizeW});
-
+      if (input.ndimension() == 3) {
+        output.resize_({sizeD, osizeH, osizeW});
+      } else {
+        output.resize_({1, sizeD, osizeH, osizeW});
+      }
       AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "adaptive_avg_pool2d_cpu", [&] {
           auto input_data = input.data<scalar_t>();
           auto output_data = output.data<scalar_t>();
@@ -260,7 +263,7 @@
     auto gradOutput = gradOutput_.contiguous();
 
     /* backprop */
-    if (input.ndimension() == 3)
+    if (input.ndimension() == 3 || input.size(-4) == 1)
     {
       AT_DISPATCH_FLOATING_TYPES_AND_HALF(
         input.scalar_type(), "adaptive_avg_pool2d_backward_cpu", [&] {