diff --git a/pytorch3d/csrc/pulsar/pytorch/tensor_util.cpp b/pytorch3d/csrc/pulsar/pytorch/tensor_util.cpp index 1dd41ed4..b372e0bf 100644 --- a/pytorch3d/csrc/pulsar/pytorch/tensor_util.cpp +++ b/pytorch3d/csrc/pulsar/pytorch/tensor_util.cpp @@ -8,6 +8,7 @@ #ifdef WITH_CUDA #include +#include #include #endif #include @@ -33,13 +34,13 @@ torch::Tensor sphere_ids_from_result_info_nograd( .contiguous(); if (forw_info.device().type() == c10::DeviceType::CUDA) { #ifdef WITH_CUDA - cudaMemcpyAsync( + C10_CUDA_CHECK(cudaMemcpyAsync( result.data_ptr(), tmp.data_ptr(), sizeof(uint32_t) * tmp.size(0) * tmp.size(1) * tmp.size(2) * tmp.size(3), cudaMemcpyDeviceToDevice, - at::cuda::getCurrentCUDAStream()); + at::cuda::getCurrentCUDAStream())); #else throw std::runtime_error( "Copy on CUDA device initiated but built " diff --git a/pytorch3d/csrc/pulsar/pytorch/util.cpp b/pytorch3d/csrc/pulsar/pytorch/util.cpp index 7d25b6e8..87eb8815 100644 --- a/pytorch3d/csrc/pulsar/pytorch/util.cpp +++ b/pytorch3d/csrc/pulsar/pytorch/util.cpp @@ -7,6 +7,7 @@ */ #ifdef WITH_CUDA +#include #include namespace pulsar { @@ -17,7 +18,8 @@ void cudaDevToDev( const void* src, const int& size, const cudaStream_t& stream) { - cudaMemcpyAsync(trg, src, size, cudaMemcpyDeviceToDevice, stream); + C10_CUDA_CHECK( + cudaMemcpyAsync(trg, src, size, cudaMemcpyDeviceToDevice, stream)); } void cudaDevToHost( @@ -25,7 +27,8 @@ void cudaDevToHost( const void* src, const int& size, const cudaStream_t& stream) { - cudaMemcpyAsync(trg, src, size, cudaMemcpyDeviceToHost, stream); + C10_CUDA_CHECK( + cudaMemcpyAsync(trg, src, size, cudaMemcpyDeviceToHost, stream)); } } // namespace pytorch