From 5e8b6f8d1000fa7df595432d48e876fc18364b6e Mon Sep 17 00:00:00 2001 From: David Hou Date: Wed, 21 Feb 2024 23:33:35 -0800 Subject: [PATCH] E501 --- tinygrad/features/multi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/features/multi.py b/tinygrad/features/multi.py index a7de9d68d9..0adb319268 100644 --- a/tinygrad/features/multi.py +++ b/tinygrad/features/multi.py @@ -101,7 +101,7 @@ class MultiLazyBuffer: # new_axis is the last one that preserves prod(prior to new_axis) and must not move items between shards new_axis = len(arg_acc) - arg_acc[::-1].index(prod(self.shape[:self.axis])) - 1 if not arg[new_axis] == self.shape[self.axis]: - assert self.shape[self.axis] % len(self.real_lbs) == 0, f"cannot reshape on-axis for uneven shards {self.axis} {self.shape} {len(self.real_lbs)}" + assert self.shape[self.axis] % len(self.real_lbs) == 0, f"cannot reshape on-axis for uneven shard {self.axis} {self.shape} {len(self.real_lbs)}" assert arg[new_axis] % len(self.real_lbs) == 0, f"new on-axis shape must divide evenly between devices {new_axis} {arg} {len(self.real_lbs)}" return MultiLazyBuffer([x.reshape(tuple(s if a != new_axis else x.shape[self.axis] if s == self.shape[self.axis] else