fix bn test, affine is True
MODIFY
test/test_nn.py
test/test_nn.py
@@ -13,7 +13,8 @@ class TestNN(unittest.TestCase):
sz = 4
# create in tinygrad
- bn = BatchNorm2D(sz, eps=1e-5, training=training, track_running_stats=training)
+ Tensor.training = training
+ bn = BatchNorm2D(sz, eps=1e-5, track_running_stats=training)
bn.weight = Tensor.randn(sz)
bn.bias = Tensor.randn(sz)
bn.running_mean = Tensor.randn(sz)

MODIFY
tinygrad/nn.py
tinygrad/nn.py
@@ -2,8 +2,8 @@ from tinygrad.tensor import Tensor
import numpy as np
class BatchNorm2D:
- def __init__(self, sz, eps=1e-5, track_running_stats=False, momentum=0.1):
- # TODO: add affine?
+ def __init__(self, sz, eps=1e-5, affine=True, track_running_stats=False, momentum=0.1):
+ assert affine == True
self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
self.weight, self.bias = Tensor.ones(sz), Tensor.zeros(sz)

fix bn test, affine is True