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 set it up (I use linux/ubuntu 18.04 and pytorch 1.5.0)
To use this 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
rm -rf build dist loss.egg-info __pycache__
pip uninstall loss -y
python setup.py install --user
python test.py

To run it 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. Let me know if this is a version problem.