Cleanup index types

pytorch · May 11, 2024 · ac6f63a · ac6f63a
1 parent 1a37190
commit ac6f63a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch/csrc/distributed/c10d/quantization/quantization_gpu.cu b/torch/csrc/distributed/c10d/quantization/quantization_gpu.cu
@@ -76,8 +76,8 @@ at::Tensor _float_to_bfloat16_cuda(const at::Tensor& input) {
     return output;
   }
 
-  constexpr int threads_per_block = 256;
-  const int blockDim_x = std::min(output_columns, threads_per_block);
+  constexpr size_t threads_per_block = 256;
+  const auto blockDim_x = std::min(output_columns, threads_per_block);
   dim3 blockDim(blockDim_x, threads_per_block / blockDim_x);
   const auto gridDim_x = (output_columns + blockDim.x - 1) / blockDim.x;
   const auto gridDim_y =