Repository: pytorch
Issue: CPU parallelization across batch has random faulty behavior on backward
Labels: triaged,module: multithreading

### 🐛 Describe the bug

Hi,

I am trying to implement a CPU loss that applies parallelization across batch.
When running the backward function I receive a very strange error:

```
Traceback (most recent call last):
  File "test.py", line 8, in <module>
    torch.sum(loss(a, b)).backward()
  File "/home/.local/lib/python3.7/site-packages/torch/_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/.local/lib/python3.7/site-packages/torch/autograd/__init__.py", line 149, in backward
    allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
RuntimeError: res[i].defined()INTERNAL ASSERT FAILED at "../torch/csrc/autograd/functions/tensor.cpp":97, please report a bug to PyTorch.
```
 
I've managed to reduce my complicated function to something simple in order to provide you with a reproducible example.

Here are the instructions on how to reproduce it:

`conda create -n pytorch-latest pytorch torchvision cpuonly -c pytorch`

Create a dir named 'loss' and add the following files below it:

setup.py
```
from setuptools import setup, Extension
from torch.utils import cpp_extension
from __init__ import __version__

cpp_module = cpp_extension.CppExtension(
    'loss',
    sources=['loss.cpp'],
    extra_compile_args=['-fopenmp'],
    extra_link_args=['-lgomp']
)

setup(name='loss',
      ext_modules=[cpp_module],
      cmdclass={'build_ext': cpp_extension.BuildExtension},
      )
```

test.py
```
import torch
torch.set_num_threads(3)

from loss import loss

a = torch.nn.Parameter(torch.ones((3, 1000)))
b = torch.zeros((3, 1000))

torch.sum(loss(a, b)).backward()
```

loss.cpp
```
#include <torch/extension.h>
#include <iostream>
#include <ATen/Parallel.h>

torch::Tensor loss(torch::Tensor z, torch::Tensor x)
{
    torch::Tensor z_out = at::empty({z.size(0), z.size(1)}, z.options());
    int64_t batch_size = z.size(0);
    at::parallel_for(0, batch_size, 0, [&](int64_t start, int64_t end) {
        for (int64_t b = start; b < end; b++)
        {
            z_out[b] = z[b] - x[b];
        }
    });
    return z_out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
    m.def("loss", &loss, "loss");
}
```

test.sh
```
conda activate pytorch-latest
rm -rf build dist loss.egg-info __pycache__
pip uninstall loss -y
python setup.py install --user
python test.py
```

> Now, move inside the 'loss' directory and run: 
> `sh test.sh`
> Note: the error doesn't always appear. You may need to run the script multiple times to see it.

### Versions

Collecting environment information...
PyTorch version: 1.10.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.4 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2
Libc version: glibc-2.27

Python version: 3.9.7 (default, Sep 16 2021, 13:09:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.4.0-19041-Microsoft-x86_64-with-glibc2.27
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.21.2
[pip3] torch==1.10.1
[pip3] torchvision==0.11.2
[conda] blas                      1.0                         mkl
[conda] cpuonly                   2.0                           0    pytorch
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-service               2.4.0            py39h7f8727e_0
[conda] mkl_fft                   1.3.1            py39hd3c417c_0
[conda] mkl_random                1.2.2            py39h51133e4_0
[conda] numpy                     1.21.2           py39h20f2e39_0
[conda] numpy-base                1.21.2           py39h79a1101_0
[conda] pytorch                   1.10.1              py3.9_cpu_0    pytorch
[conda] pytorch-mutex             1.0                         cpu    pytorch
[conda] torchvision               0.11.2                 py39_cpu  [cpuonly]  pytorch