From 6da6b52b4cf9c91d90e35ecca7126f3b28dca813 Mon Sep 17 00:00:00 2001 From: ghostplant Date: Tue, 9 Apr 2024 15:58:08 +0800 Subject: [PATCH] replace unnecessary zeros -> empty (#229) --- tutel/impls/communicate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tutel/impls/communicate.py b/tutel/impls/communicate.py index 9b9543d..51ba737 100644 --- a/tutel/impls/communicate.py +++ b/tutel/impls/communicate.py @@ -220,7 +220,7 @@ def batch_all_to_all_v(datas, partition_sizes, group=None): return list(datas), in_sizes out_sizes = simple_all_to_all(in_sizes, group=group) datas = [data.contiguous().view(-1).cuda() for data in datas] - outputs = [torch.zeros([out_sizes.sum()], dtype=data.dtype, device=data.device) for data in datas] + outputs = [torch.empty([out_sizes.sum()], dtype=data.dtype, device=data.device) for data in datas] tutel_custom_kernel.batch_all_to_all_v(datas, outputs, in_sizes, out_sizes) return outputs, out_sizes @@ -234,7 +234,7 @@ def batch_all_gather_v(datas, group=None): return list(datas), input_size output_sizes = simple_all_gather(input_size) size_int = int(output_sizes.sum()) - outputs = [torch.zeros([size_int], dtype=data.dtype, device=data.device) for i, data in enumerate(datas)] + outputs = [torch.empty([size_int], dtype=data.dtype, device=data.device) for i, data in enumerate(datas)] tutel_custom_kernel.batch_all_gather_v(datas, outputs, output_sizes) return outputs, output_sizes