From e5ba385f03ed22af5ab04553ef1fd8919ee5a736 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 3 Jul 2024 19:42:56 -0400 Subject: [PATCH] remove first contiguous in multi from_sharded (#5121) second contiguous guarantees lbs are contiguous going into MultiLazyBuffer, don't need the first contiguous --- tinygrad/multi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/multi.py b/tinygrad/multi.py index 848cbe4bd8..7363c14cbc 100644 --- a/tinygrad/multi.py +++ b/tinygrad/multi.py @@ -71,7 +71,7 @@ class MultiLazyBuffer: @staticmethod def from_sharded(lb:LazyBuffer, devices:Tuple[str, ...], axis:Optional[int]=None): - lbs = [lb.contiguous() if lb.base != lb and not lb.is_unrealized_unmasked_const() else lb] * len(devices) + lbs = [lb] * len(devices) sharded_lbs = [lb.copy_to_device(d) for lb,d in zip(to_sharded(lbs, axis) if axis is not None else lbs, devices)] return MultiLazyBuffer([lb if lb.is_unrealized_unmasked_const() else lb.contiguous(allow_buffer_view=False) for lb in sharded_lbs], axis)