[ROCm] Update CUDAPluggableAllocator.h (#1984) (pytorch#153974)

[ROCm] Update CUDAPluggableAllocator.h (#1984) (pytorch#150010)

Altering the flag to use the correct streamType in CUDAPluggableAllocator class for ROCm gpu. The flag TORCH_HIP_VERSION does not work for ROCm as intended. This flag is replaced with USE_ROCM. This is impacting Distributed Fused Adam in Rocm/APEX when using nccl_ub feature. This has been tested with rocm/apex.

See PR ROCm/apex#184

Pull Request resolved: pytorch#150010
Approved by: https://github.com/jeffdaily

(cherry picked from commit a19b667)

Co-authored-by: Sriram Kumar <[email protected]>

Author: pytorchbot

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;