From 6cac3c99e2cb1f41953bdcace7bb74c1b3fb3d36 Mon Sep 17 00:00:00 2001 From: David Hou Date: Wed, 21 Feb 2024 23:34:06 -0800 Subject: [PATCH] argfix axis --- tinygrad/tensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 810b3577f8..90bcbf2df9 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -928,7 +928,7 @@ class Tensor: shape = tuple(s if ax in argfix(axis) else 1 for ax, s in enumerate(self.shape)) x = self - mean.reshape(shape) if weight: x = x * weight.reshape(shape) - ret = x.mul(invstd.reshape(shape) if len(invstd.shape) == len(axis) else invstd) + ret = x.mul(invstd.reshape(shape) if len(invstd.shape) == len(argfix(axis)) else invstd) return (ret + bias.reshape(shape)) if bias else ret def dropout(self, p=0.5) -> Tensor: