Fix: torch::kBFloat16

2026-05-14 10:21:10 +00:00
parent 3f62e49e6e
commit a272bc49b0
1 changed files with 1 additions and 1 deletions
--- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp
+++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp
@@ -20,7 +20,7 @@ torch::Tensor cutlass_nvfp4_gemm_forward(
    torch::Tensor SFB,
    int64_t M, int64_t N, int64_t K
 ) {
-    auto D = torch::empty({M, N}, torch::dtype(at::kBF16).device(A_packed.device()));
+    auto D = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A_packed.device()));

    auto stream = c10::cuda::getCurrentCUDAStream();
    cudaStream_t cuda_stream = stream.stream();