From ebd7290f561afd45752f74fc6cb84ec96a601a97 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sun, 5 Jun 2022 16:51:51 -0700 Subject: [PATCH] simpler batchnorm --- tinygrad/nn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/nn.py b/tinygrad/nn.py index e39fecd1b0..f251fe76df 100644 --- a/tinygrad/nn.py +++ b/tinygrad/nn.py @@ -30,7 +30,7 @@ class BatchNorm2D: def normalize(self, x, mean, var): x = (x - mean.reshape(shape=[1, -1, 1, 1])) * self.weight.reshape(shape=[1, -1, 1, 1]) - return x.div(var.add(self.eps).reshape(shape=[1, -1, 1, 1])**0.5) + self.bias.reshape(shape=[1, -1, 1, 1]) + return x.mul(var.add(self.eps).reshape(shape=[1, -1, 1, 1])**-0.5) + self.bias.reshape(shape=[1, -1, 1, 1]) class Conv2d: def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, bias=True):