From 5743a0cc85f7de5c8d42b14d28bc1941dfeb875c Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Tue, 25 Mar 2025 17:09:01 +0200 Subject: [PATCH] [release/2.5] Update CUDAPluggableAllocator.h Altering the flag to use the correct streamType for CUDAPluggableAllocator. This is impacting Distributed Fused Adam in Rocm/APEX. Apex Commit : https://github.com/ROCm/apex/pull/188 --- torch/csrc/cuda/CUDAPluggableAllocator.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/csrc/cuda/CUDAPluggableAllocator.h b/torch/csrc/cuda/CUDAPluggableAllocator.h index 8652ef0f2bfde8..4827d206d6952f 100644 --- a/torch/csrc/cuda/CUDAPluggableAllocator.h +++ b/torch/csrc/cuda/CUDAPluggableAllocator.h @@ -37,7 +37,7 @@ struct TORCH_CUDA_CPP_API CUDAPluggableAllocatorDeleterContext { cudaStream_t stream_; }; -#if defined(TORCH_HIP_VERSION) +#if defined(USE_ROCM) using streamType = c10::hip::HIPStream; #else using streamType = c10::cuda::CUDAStream;