[DeprecationWarning] Replacing soon-to-be-deprecated lazyInitCUDA w…

…ith `lazyInitDevice(at::kCUDA)` (#7855)
dmlc · Jan 16, 2025 · ba73133 · ba73133
1 parent 17017c2
commit ba73133
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/graphbolt/src/cuda/common.h b/graphbolt/src/cuda/common.h
@@ -58,7 +58,9 @@ struct CUDAWorkspaceAllocator {
   // Required by thrust to satisfy allocator requirements.
   using value_type = value_t;
 
-  explicit CUDAWorkspaceAllocator() { at::globalContext().lazyInitCUDA(); }
+  explicit CUDAWorkspaceAllocator() {
+    at::globalContext().lazyInitDevice(at::kCUDA);
+  }
 
   template <class U>
   CUDAWorkspaceAllocator(CUDAWorkspaceAllocator<U> const&) noexcept {}

diff --git a/tensoradapter/pytorch/torch.cpp b/tensoradapter/pytorch/torch.cpp
@@ -28,7 +28,7 @@ TA_EXPORTS void CPURawDelete(void* ptr) {
 
 #ifdef DGL_USE_CUDA
 TA_EXPORTS void* CUDARawAlloc(size_t nbytes, cudaStream_t stream) {
-  at::globalContext().lazyInitCUDA();
+  at::globalContext().lazyInitDevice(at::kCUDA);
   return c10::cuda::CUDACachingAllocator::raw_alloc_with_stream(nbytes, stream);
 }