ATen/cuda/Atomic.cuh does not exist before Pytorch 1.11. Therefore, this PR is to make index_mul_2d extension compatible with various Pytorch releases.
New Pytorch releases (after 1.11): #include <ATen/cuda/Atomic.cuh>
Old Pytorch releases (before 1.11): #include <THC/THCAtomics.cuh>
ATen/cuda/Atomic.cuh does not exist before Pytorch 1.11. Therefore, this PR is to make index_mul_2d extension compatible with various Pytorch releases.