diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 55ee4a69c1..488fc0774c 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -324,7 +324,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): def detach(self): return UOp(Ops.DETACH, self.dtype, (self,)) def index(self, *srcs:UOp|None, **kwargs): return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype), (self,)+tuple([x for x in srcs if x is not None]), **kwargs) - def __getitem__(self, idx): return self.index(idx) + def __getitem__(self, *idx): return self.index(*idx) def const_like(self, b:ConstLike): # constants can optionally have a DEVICE source return UOp.const(self.dtype, b, device=self._device, shape=self._shape)