Skip to content

Commit a2c9905

Browse files
committed
use add_(other, alpha) instead of explicity add(other.mul(a))
1 parent 78c9f56 commit a2c9905

File tree

1 file changed

+6
-8
lines changed

1 file changed

+6
-8
lines changed

bluefog/torch/mpi_ops.cc

+6-8
Original file line numberDiff line numberDiff line change
@@ -404,12 +404,11 @@ int DoNeighborAllreduce(::torch::Tensor tensor, ::torch::Tensor output,
404404
output_reduced.mul_(weight);
405405
} else {
406406
output_reduced.add_(
407-
output_buffer.slice(0, i * first_dim, (i + 1) * first_dim)
408-
.mul_(weight));
407+
output_buffer.slice(0, i * first_dim, (i + 1) * first_dim), weight);
409408
}
410409
}
411410
output_buffer.resize_(shape_vector);
412-
output_buffer.add_(tensor_buffer.mul(self_weight));
411+
output_buffer.add_(tensor_buffer, self_weight);
413412
if (is_hierarchical){
414413
// Because there is ncclAllreduce just take sum.
415414
output_buffer.div_(bluefog_local_size());
@@ -495,12 +494,11 @@ int DoNeighborAllreduce(::torch::Tensor tensor, ::torch::Tensor output,
495494
output_reduced.mul_(weight);
496495
} else {
497496
output_reduced.add_(
498-
output_buffer.slice(0, i * first_dim, (i + 1) * first_dim)
499-
.mul_(weight));
497+
output_buffer.slice(0, i * first_dim, (i + 1) * first_dim), weight);
500498
}
501499
}
502500
output_buffer.resize_(shape_vector);
503-
output_buffer.add_(tensor_buffer.mul(self_weight));
501+
output_buffer.add_(tensor_buffer, self_weight);
504502
if (is_hierarchical){
505503
// Because there is ncclAllreduce just take sum.
506504
output_buffer.div_(bluefog_local_size());
@@ -576,7 +574,7 @@ int DoPairGossip(::torch::Tensor tensor, ::torch::Tensor output,
576574
output_buffer.add_(tensor_buffer).div_(2);
577575
} else {
578576
output_buffer.mul_(pair_weight)
579-
.add_(tensor_buffer.mul(self_weight));
577+
.add_(tensor_buffer, self_weight);
580578
}
581579
MaybeCopyBufferBack(output, output_buffer);
582580
}));
@@ -597,7 +595,7 @@ int DoPairGossip(::torch::Tensor tensor, ::torch::Tensor output,
597595
output_buffer.add_(tensor_buffer).div_(2);
598596
} else {
599597
output_buffer.mul_(pair_weight)
600-
.add_(tensor_buffer.mul(self_weight));
598+
.add_(tensor_buffer, self_weight);
601599
}
602600
MaybeCopyBufferBack(output, output_buffer);
603601
}));

0 commit comments

Comments
 (0)