Repository: pytorch
Issue: input.dim() == 4 INTERNAL ASSERT FAILED mkldnn/Pooling.cpp:201
Labels: high priority,needs reproduction,module: nn,triaged,module: mkldnn,module: pooling

### 🐛 Describe the bug

> input.dim() == 4 INTERNAL ASSERT FAILED at /pytorch/aten/src/ATen/native/mkldnn/Pooling.cpp:201, please report a bug to PyTorch. mkldnn_adaptive_avg_pool2d: Expect 2D input
> Traceback (most recent call last):
>   File "main.py", line 130, in handler
>     ocr_result_content = process_page_file_with_ocr(layout_result_content, [image])
>   File "/ocr-analysis/ocr_page.py", line 72, in process_page_file_with_ocr
>     ocr_results = get_ocr_predictions(line_and_idx_and_img_and_txt)
>   File "/ocr-analysis/ocr_page.py", line 11, in get_ocr_predictions
>     text_results, confidences = ocr_predictor.predict(images)
>   File "/ocr-analysis/ocr.py", line 292, in predict
>     test_predictions = self.run_inference(test_loader)
>   File "/ocr-analysis/ocr.py", line 303, in run_inference
>     inference_result += self.inference_one_batch(batch)
>   File "/ocr-analysis/ocr.py", line 308, in inference_one_batch
>     outputs = self.model(image)
>   File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
>     result = self.forward(*input, **kwargs)
>   File "/ocr-analysis/ocr.py", line 149, in forward
>     x = self.avg_pool(x)
>   File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
>     result = self.forward(*input, **kwargs)
>   File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/pooling.py", line 1031, in forward
>     return F.adaptive_avg_pool2d(input, self.output_size)
>   File "/usr/local/lib/python3.8/site-packages/torch/nn/functional.py", line 768, in adaptive_avg_pool2d
>     return torch._C._nn.adaptive_avg_pool2d(input, _output_size)
> RuntimeError: input.dim() == 4 INTERNAL ASSERT FAILED at /pytorch/aten/src/ATen/native/mkldnn/Pooling.cpp:201, please report a bug to PyTorch. mkldnn_adaptive_avg_pool2d: Expect 2D input

The model looks like this

```
class RecognitionModel(nn.Module):
    def __init__(self, feature_extractor, time_feature_count, lstm_hidden, lstm_len, n_class):
        super(RecognitionModel, self).__init__()
        self.feature_extractor = feature_extractor
        self.avg_pool = nn.AdaptiveAvgPool2d(
            (time_feature_count, time_feature_count))
        self.bilstm = BiLSTM(time_feature_count, lstm_hidden, lstm_len)
        self.classifier = nn.Sequential(
            nn.Linear(lstm_hidden * 2, time_feature_count),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(time_feature_count, n_class)
        )

    def forward(self, x):
        x = self.feature_extractor(x)
        b, c, h, w = x.size()
        x = x.reshape(b, c * h, w)
        x = self.avg_pool(x)
        x = x.transpose(1, 2)
        x = self.bilstm(x)
        x = self.classifier(x)
        return x
```
The model worked before I tried to switch to mkldnn.


### Versions


> Collecting environment information...
> PyTorch version: 1.4.0
> Is debug build: False
> CUDA used to build PyTorch: 10.1
> ROCM used to build PyTorch: N/A
> 
> OS: Debian GNU/Linux 10 (buster) (x86_64)
> GCC version: (Debian 8.3.0-6) 8.3.0
> Clang version: Could not collect
> CMake version: Could not collect
> Libc version: glibc-2.28
> 
> Python version: 3.8.3 (default, Jun  9 2020, 17:39:39)  [GCC 8.3.0] (64-bit runtime)
> Python platform: Linux-5.4.0-42-generic-x86_64-with-glibc2.2.5
> Is CUDA available: False
> CUDA runtime version: No CUDA
> GPU models and configuration: No CUDA
> Nvidia driver version: No CUDA
> cuDNN version: No CUDA
> HIP runtime version: N/A
> MIOpen runtime version: N/A
> 
> Versions of relevant libraries:
> [pip3] numpy==1.22.2
> [pip3] torch==1.4.0
> [pip3] torchvision==0.5.0
> [conda] Could not collect

cc @ezyang @gchanan @zou3519 @albanD @mruberry @jbschlosser @walterddr @kshitij12345 @gujinghui @PenghuiCheng @XiaobingSuper @jianyuh @VitalyFedyunin