diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py
index 6490e69a60..403f475702 100644
--- a/vllm/model_executor/determinism/batch_invariant.py
+++ b/vllm/model_executor/determinism/batch_invariant.py
@@ -918,7 +918,7 @@ def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None =
         for d in sorted_dims:
             result = result.squeeze(d)
 
-    return result
+    return result if dtype is not None else result.to(input.dtype)
 
 
 @triton.jit