diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index c3058875d1..ac5cd2de92 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -224,7 +224,7 @@ class Tensor: def reshape(self, shape, *args) -> Tensor: new_shape = argfix(shape, *args) - assert all(x != 0 for x in new_shape), f"zeros not allowed in shape {new_shape}" + assert 0 not in new_shape, f"zeros not allowed in shape {new_shape}" return mlops.Reshape.apply(self, shape=tuple(-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape)) def expand(self, shape, *args) -> Tensor: return mlops.Expand.apply(self, shape=tuple(x if x != -1 else s for s,x in zip(self.shape, argfix(shape, *args)))) def permute(self, order, *args) -> Tensor: return mlops.Permute.apply(self, order=argfix(order, *args))