Address review's comment

youngeunkwon0405 · web-flow · commit 8a62df3f310e · 2026-01-30T15:40:38.000-08:00
diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
@@ -721,7 +721,7 @@ def _does_param_require_new_bucket(param):
             # Since nccl communicator group is created lazily, we need to perform a warmup call to
             # initialize NCCL comm buffers for this dp_group before doing buffer registration.
             torch.distributed.barrier()
-            tmp_warmup_tensor = torch.tensor([0.0], device="cuda")
+            tmp_warmup_tensor = torch.zeros([1], device="cuda")
             torch.distributed.all_reduce(tmp_warmup_tensor, group=self.data_parallel_group)
             torch.distributed.barrier()
         else:
