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: https://github.com/sunset1995/torch_efficient_distloss/blob/03ea697ca261e839cae10e2ccc25cb350c04d226/torch_efficient_distloss/cuda/segment_cumsum_kernel.cu#L1 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.