[ROCm6.4_internal_testing] Update CUDAPluggableAllocator.h (#1984)

amd-sriram · amd-sriram · commit ea4e08c2d956 · 2025-03-26T08:19:45.000Z
Altering the flag to use the correct streamType for CUDAPluggableAllocator. This is impacting Distributed Fused Adam in Rocm/APEX. See PR ROCm/apex#184 Related Issue : https://ontrack-internal.amd.com/browse/SWDEV-519796
diff --git 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;
