using stream functions from at instead of c10, small fix

NVIDIA · Apr 4, 2024 · 57a5608 · 57a5608
1 parent c25a0fe
commit 57a5608
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/torch_harmonics/csrc/disco/disco_cuda_bwd.cu b/torch_harmonics/csrc/disco/disco_cuda_bwd.cu
@@ -286,7 +286,7 @@ torch::Tensor disco_cuda_bwd(torch::Tensor inp,
   torch::Tensor out = torch::zeros(out_dims, options);
 
   // get stream
-  auto stream = at::cuda::getCurrentCUDAStream();
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
 
   // assert
   static_assert(0 == (ELXTH_MAX%2));

diff --git a/torch_harmonics/csrc/disco/disco_cuda_fwd.cu b/torch_harmonics/csrc/disco/disco_cuda_fwd.cu
@@ -264,7 +264,7 @@ torch::Tensor disco_cuda_fwd(torch::Tensor inp,
   torch::Tensor out = torch::zeros(out_dims, options);
 
   // get stream
-  auto stream = at::cuda::getCurrentCUDAStream();
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
 
   // assert
   static_assert(0 == (ELXTH_MAX%2));