diff --git a/csrc/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu
index cd80bfda7dfd..99c52ef17d08 100644
--- a/csrc/moe/topk_softmax_kernels.cu
+++ b/csrc/moe/topk_softmax_kernels.cu
@@ -573,7 +573,7 @@ void topk_softmax(
             stream);
     }
     else {
-        TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
+        assert(topk_indices.scalar_type() == at::ScalarType::Int64);
         vllm::moe::topkGatingSoftmaxKernelLauncher(
             gating_output.data_ptr<float>(),
             topk_weights.data_ptr<float>(),