Use .scalar_type() in AT_DISPATCH

As `AT_DISPATCH_FLOATING_TYPES(x.type()` pattern is deprecated since pytorch/pytorch#17996
lab-cosmo · Nov 7, 2024 · b003104 · b003104
1 parent 877370a
commit b003104
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/sphericart-torch/src/torch_cuda_wrapper.cpp b/sphericart-torch/src/torch_cuda_wrapper.cpp
@@ -9,7 +9,7 @@
 #include "cuda_base.hpp"
 #include "sphericart.hpp"
 
-#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
 #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
 #define CHECK_SAME_DTYPE(x, y)                                                                     \
     TORCH_CHECK(x.scalar_type() == y.scalar_type(), #x " and " #y " must have the same dtype.")
@@ -34,7 +34,7 @@ torch::Tensor sphericart_torch::spherical_harmonics_backward_cuda(
         xyz_grad = torch::empty_like(xyz);
 
         AT_DISPATCH_FLOATING_TYPES(
-            xyz.type(), "spherical_harmonics_backward_cuda", ([&] {
+            xyz.scalar_type(), "spherical_harmonics_backward_cuda", ([&] {
                 sphericart::cuda::spherical_harmonics_backward_cuda_base<scalar_t>(
                     dsph.data_ptr<scalar_t>(),
                     sph_grad.data_ptr<scalar_t>(),
@@ -48,4 +48,4 @@ torch::Tensor sphericart_torch::spherical_harmonics_backward_cuda(
     }
     // synchronization happens within spherical_harmonics_backward_cuda_base
     return xyz_grad;
-}
+}