From 4c3895744edede50c71f73aff74b32cfd073f2a4 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 4 Oct 2024 09:03:56 -0400 Subject: [PATCH] type annotation for layernorm (#6883) --- tinygrad/tensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index feb9b33824..0eb8345c6f 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -3042,7 +3042,7 @@ class Tensor: """ return functools.reduce(lambda x,f: f(x), ll, self) - def layernorm(self, axis=-1, eps:float=1e-5) -> Tensor: + def layernorm(self, axis:Union[int,Tuple[int,...]]=-1, eps:float=1e-5) -> Tensor: """ Applies Layer Normalization over a mini-batch of inputs.