update base docker image

ROCm · Jun 5, 2024 · a6af475 · a6af475
1 parent 69ce080
commit a6af475
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/Dockerfile.rocm b/Dockerfile.rocm
@@ -1,5 +1,5 @@
 # default base image
-ARG BASE_IMAGE="rocm/pytorch:rocm6.1.1_ubuntu20.04_py3.9_pytorch_release-2.1.2"
+ARG BASE_IMAGE="rocm/pytorch:rocm6.1.1_ubuntu20.04_py3.9_pytorch_staging"
 
 ARG COMMON_WORKDIR=/app
 

diff --git a/csrc/quantization/fp8/gemm_kernel.cu b/csrc/quantization/fp8/gemm_kernel.cu
@@ -101,7 +101,7 @@ torch::Tensor fp8_gemm(torch::Tensor& a, torch::Tensor& b, torch::Tensor& scaleA
     auto d_scaleD = scaleD.data_ptr();
 
     auto handle = at::cuda::getCurrentCUDABlasLtHandle();
-    auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
+    auto stream = at::cuda::getCurrentCUDAStream();
 
     hipblaslt_ext::GemmPreference gemmPref;
     gemmPref.setMaxWorkspaceBytes(0);
@@ -218,7 +218,7 @@ torch::Tensor fp8_gemm_16(
     auto d_scaleB = transpose_result ? scaleA.data_ptr() : scaleB.data_ptr();
 
     auto handle = at::cuda::getCurrentCUDABlasLtHandle();
-    auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
+    auto stream = at::cuda::getCurrentCUDAStream();
 
     hipblaslt_ext::GemmPreference gemmPref;
     gemmPref.setMaxWorkspaceBytes(0);