Merge pull request #35 from ROCm/charlifu/update_base_image_with_newe…

…r_pytorch Update base docker image with Pytorch 2.3
ROCm · Jun 5, 2024 · 68cdb95 · 68cdb95
2 parents 69ce080 + a6af475
commit 68cdb95
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);