Replace <torch/extension.h> with <torch/all.h>

[tobias-kirschstein]

Thank you for this great implementation of distortion loss.

I ran into this issue when compiling your extension.

\lib\site-packages\torch\include\pybind11\cast.h(624): error: too few arguments for template template parameter "Tuple"
          detected during instantiation of class "pybind11::detail::tuple_caster<Tuple, Ts...> [with Tuple=std::pair, Ts=<T1, T2>]"
(721): here

\lib\site-packages\torch\include\pybind11\cast.h(717): error: too few arguments for template template parameter "Tuple"
          detected during instantiation of class "pybind11::detail::tuple_caster<Tuple, Ts...> [with Tuple=std::pair, Ts=<T1, T2>]"
(721): here

Following a discussion in the Pytorch3D repo it seems like the issue is that nvcc doesn't like pybind under Windows.
I found out that replacing this line:

torch_efficient_distloss/torch_efficient_distloss/cuda/segment_cumsum_kernel.cu

Line 1 in 03ea697

#include <torch/extension.h>

with

#include <torch/all.h> 

actually solved the problem for me.

If there is no other reason to use torch/extension.h there, I would suggest to change that line in your repo to facilitate compilation for some OS-nvcc pairs.